¹

²

³

⁴

1 Introduction↩︎

Deep neural networks (DNNs) have emerged as a central technology to extract statistical regularities from data at unprecedented scale. The core promise of this machine learning technique is that, aside from a computational speedup compared to classical kernel methods, DNNs can adapt their internal feature representation during training, and thereby approximate a kernel method where the kernel has been automatically adapted for each task (for good or worse). However, the physical laws governing this kernel adaptation, and therefore the precise capabilities and limitations of this method, have turned out to be highly resistant to theoretical inquiry for networks that are both multi-layered and nonlinear.

DNNs are objects well suited for study via statistical physics techniques, because the network size, the dimensionality of input data, and the number of data examples are all simultaneously large. Knowing the physical laws of the system then amounts to finding an effective description in terms of low-dimensional order parameters which capture the system behavior. What is less obvious, is the best large system limit in which to study deep networks, that is the relative scaling of layer sizes, initialization parameters, training procedures, and data set sizes to infinity.

The large variety of choices, affecting both conclusions and tractability, can be significantly reduced by adopting the Bayesian viewpoint of asking for the posterior given a prior over parameters closely related to the initialization scheme. In this way, the need to choose a training algorithm and its hyperparameter scalings is removed. Furthermore, while algorithmic biases constitute an important field of research [1]–[5], it has been shown that such biases may quantitatively but not qualitatively alter the performance compared to realizations drawn from the Bayesian network posterior [6].

We can then sketch the landscape of limits for Bayesian networks as follows: In the Neural Network Gaussian Process (NNGP) limit, also termed "lazy" infinite-width limit, where all layer widths are sent to infinity while the dataset is fixed, the network becomes equivalent to Gaussian process regression with a fixed NNGP kernel determined by the architecture [7]–[11]. The lazy regime was contrasted by another, "rich" infinite width limit, where the network output in the prior or at initialization is scaled down to vanish with width, forcing a strong adaptation of the network features to fit the \(O(1)\) target labels [12]–[15]. However, while realistic deep networks are large, their width is typically not large compared to the number of data samples as assumed by these limits. Indeed, it was proposed that finite-width effects are responsible for the advantages of DNNs over kernel methods [16]. Another way to capture these effects is the proportional limit, where layer widths, input dimension, and sample size scale proportionally to infinity. This limit is classical in the statistical physics of shallow (at most one hidden layer) networks [17]–[20]. For DNNs, which have a number of parameters quadratic in the hidden-layer width, also supraproportional limits with faster sample-size scaling could in some cases be relevant [21], [22]. These are unexplored and expected to be extremely difficult to study for non-random data.

In this work, we focus on finite-width, multi-layer and nonlinear DNNs, adopting the proportional limit, and propose an effective low-dimensional, albeit approximate, theory. Considerable theoretical understanding exists when any one of these three ingredients is removed: nonlinearity, as in deep linear networks [23]–[28]; depth, as in single-index models [29]–[32] and shallow networks [12], [13], [33]–[35]; or finite-width effects, as in the infinite-width NNGP limit [7]–[9], [36]–[38]. The simultaneous presence of all three, arguably central for deep learning, so far evaded low-dimensional description however. By choosing an approximate, Ansatz based theory, we trade off asymptotic exactness for practical insight into nonlinear DNN generalization when trained on real datasets.

1.1 The proportional regime of Bayesian DNNs: previous results and open problems↩︎

The proportional limit in Bayesian deep linear neural networks was first considered in Ref. [23]. Here, the authors introduced a method, the Backpropagating kernel renormalization, which allows for the incremental integration of the network weights layer by layer starting from the network output layer and progressing backward until the first layer’s weights are integrated out. Using this approach, properties of the network at fixed \(\alpha = P/N\) are evaluated via a self-consistent equation for a one-dimensional order parameter.

Afterwards, it was shown that the validity of this approach in the proportional regime is a consequence of an exact representation for the partition function of deep linear networks at finite \(N\) and \(P\). Such a representation was first given in terms of Meijer-G functions for the case of fully-connected architectures with a single readout neuron [26], [39], and later generalized to the case of networks with multiple outputs and convolutional layers [24]. Here the authors noted that the joint prior distribution of the outputs for a finite deep linear architecture with \(L\) hidden layers can be expressed as a mixture of Gaussians over \(L\) low-dimensional positive-definite matrix ensembles. This exact integral representation for the prior directly transfers to the partition function in the case of Gaussian likelihood (square-loss), and it provides a mathematically rigorous framework for the Backpropagating kernel renormalization method introduced in [23].

In Ref. [23], the authors also proposed a heuristic extension of the theory to fully-connected DNNs with ReLU activation, based on the observation that \(\mathrm{ReLU}(\gamma x) = \gamma \mathrm{ReLU}(x)\), \(\forall \, \gamma > 0\). Surprisingly, the resulting theory, which is obtained replacing the trivial NNGP linear kernel with the ReLU one, predicts the generalization performance of finite DNNs with ReLU activation for modest number of training patterns \(P \sim 10^2\) and depth \(L < 5\).

The work in Ref. [40] showed that a variational low-dimensional free-energy (or effective action) emerges in non-linear one-hidden layer networks (\(L = 1\)) with generic activation function, thanks to a Gaussian equivalence informally justified using Breuer-Major theorems [41], [42]. This framework was successfully generalized later to investigate 1HL convolutional architectures [43], FC networks with multiple outputs [44], [45], transfer and continual learning [46], [47]. In all these cases, new interesting forms of kernel shape renormalization arise, at variance with the 1HL fully-connected (and single output) case, where the renormalized kernel is found as a global, data-dependent rescaling of the NNGP kernel.

Crucially, neither Ref. [23] nor Ref. [40] elaborate convincing and principled arguments to establish whether an equivalent dimensional reduction of the partition function to the one found for \(L = 1\) holds for deeper (\(L > 1\)) architectures with generic activation function.

Another line of work seeking to explain feature learning in DNNs focuses not solely on going beyond the infinite width limit in the standard parametrization, but also on using a different scaling of the network output with \(N\), termed the mean-field [12]–[14] or \(\mu\)P [15] parametrization, which has become popular in the theoretical literature because even in the infinite-width limit the activations in all layers change by \(O(1)\) during training. This has been termed the rich learning regime as the empirical kernels strongly differ between initialization and trained network, and between prior and posterior, indicating adaptation to the data. The empirical kernels here are defined as the Gram matrices of activations in a layer, and therefore encode the physically learned representation [48], [49]. They are conceptually distinct from the effective kernels arising in the kernel shape renormalization approach, coinciding only in the NNGP limit, but both can be used to predict network outputs, as seen in linear networks where exact descriptions of both objects are available [23], [24], [28]. A series of ground-breaking papers has been seeking to predict the adaptation of the empirical kernels in nonlinear networks, yet by keeping the full \(P\times P\) kernel matrices at each layer as high-dimensional observables [50]–[53]. For shallow \(L=1\) networks, \(P\)-dimensional order parameters have been considered [28], [53]. In the mean-field setting, the empirical kernels concentrate in the infinite-width limit to nontrivial posterior configurations [50], [52], [54]. In the standard parametrization, the more realistic proportional regime yields adaptation of the empirical kernels through finite-width fluctuations [51]. Overall, this approach is successful but recasts training and inference on a task in a coupled set of \(P \times P\) dimensional self-consistency equations, which does not correspond to a dimensionality reduction with respect to the \(N \times N\) parameter space in the proportional regime, leads to significant numerical difficulty in solving the self-consistency equations, and demonstrates the need for an effective low-dimensional theory for nonlinear DNNs.

1.2 Summary of the major results↩︎

In this work, we study the statistical mechanics of learning in overparameterized Bayesian DNNs with general nonlinearities in the hidden units, providing a low-dimensional approximate description of such architectures at finite width, thus going beyond the NNGP large-width limit. In particular, we make the following contributions:
(i) We introduce the Equivalent Wishart Ansatz (EWA) in the proportional regime (Sec. 4), which provides an approximate description of the fluctuations of the hierarchical empirical kernels in nonlinear DNNs. Leveraging this novel framework, we obtain a low-dimensional integral representation of the characteristic function of the prior for generic \(L\)-layer NNs with fully connected (both scalar and multiple outputs) and convolutional layers (Sec. 4, 5.1 and 5.2). We derive the corresponding Large Deviation Principle for the emerging \(L\) order parameters, and test the validity of the EWA by comparing the theoretical rate functions with empirical simulations sampling the ground-truth MLP prior, showing convergence close to the asymptotic predicted behavior (Sec. 4.2).
(ii) We identify different deep kernel renormalization schemes at the level of the prior, that encode distinct representation learning capabilities in deep architectures. In addition, we provide a low-dimensional expression for the posterior partition function in terms of an effective action (Sec. 4.3), which explicitly depends only on the order parameters and the dataset, and non-perturbatively encodes the effect of depth and width. We conducted an extensive simulation campaign to assess the predictive power of the aforementioned theory, comparing the analytical learning curves obtained with the Equivalent Wishart Ansatz for the predictor statistics against Bayesian learning experiments employing finite networks (Sec. 6.3). Overall, we find that the theory captures the empirical behavior at finite width up to \(L\sim O(10)\) and for different datasets, activation functions, and numbers of training examples. We also discuss the emergence of discrepancies when both the depth and the dataset size are simultaneously large (Sec. 6.4), in particular describing a so far unreported metastability phenomenon occurring in the posterior at moderate depths \(L>5\).
(iii) Using both the standard (SP) and mean-field (\(\mu\)P) parametrizations, we show that in the \(P \sim N\) regime covered, a collective macroscopic description of feature learning is possible, without introducing \(P \times P\) dimensional self-consistency equations (Sec. 6.6). In the EWA, improvements in generalization in the mean-field regime are shown to be driven by a suppression of the predictor variance as a function of the width, rather than by strong adaptation of the predictor bias. Overall, these Ansatz-based and empirically successful results pave the way for a low-dimensional and asymptotically exact theory of kernel adaptation in DNNs in the proportional regime.

2 Setting of the problem and notations↩︎

We consider deep neural networks with \(L\) fully-connected hidden layers, where the pre-activations of each layer \(h_{i_{\ell}}^{(\ell)}\) (\(i_{\ell} = 1,\dots, N_{\ell}\); \(\ell = 1, \dots, L\)) are given recursively as a non-linear function of the pre-activations at the previous layer \(h_{i_{\ell-1}}^{(\ell-1)}\) (\(i_{\ell-1}= 1, \dots, N_{\ell-1}\)): \[\begin{align} \begin{aligned} h_{i_\ell}^{(\ell)}(x) &= \frac{1}{\sqrt {N_{\ell-1}}} \sum_{i_{\ell-1}=1}^{N_{\ell-1}} W^{(\ell)}_{i_{\ell}i_{\ell-1}} \sigma\!\left(h_{i_{\ell-1}}^{(\ell-1)}(x)\right) + b_{i_\ell}^{(\ell)}\,, \label{eq:preactivations-definition}\\ h_{i_1}^{(1)}(x) &= \frac{1}{\sqrt {N_{0}}} \sum_{i_0=1}^{N_{0}} W^{(1)}_{i_1 i_0} x_{i_0} + b_{i_1}^{(1)}\;, \end{aligned} \end{align}\tag{1}\] where \(W^{(\ell)}\) and \(b^{(\ell)}\) are respectively the weights and the biases of the \(\ell\)-th layer, whereas the input layer has dimension \(N_0\) (the input data dimension). \(\sigma\) is a non-linear activation function and it is common to each layer. For brevity, in the calculations we will restrict to the case without biases, \(b^{(\ell)}=0\). We add one last readout layer and define the function implemented by the deep neural network as: \[\label{eq:f95DNN} f_{\theta} ( x) = \frac{1}{\sqrt{N_L}} \sum_{i_L=1}^{N_L} W^{(L+1)}_{i_L} \sigma \!\left[ h_{i_L}^{(L)} (x)\right]\,,\tag{2}\] where \(W^{(L+1)}\) is the vector of weights of the last layer and \(\theta\) indicates the collection of all the weights of the network, \(\theta = \{ W^{(\ell)}\}_\ell\).

We consider a supervised learning problem with fixed training set \(\mathcal{D}_P = \{X,Y\}=\{x^\mu, y^\mu\}_{\mu=1}^P\), where each \(x^\mu \in \mathbb{R}^{N_0}\) and the corresponding labels \(y^\mu \in \mathbb{R}\). We analyse regression problems with squared-error loss function: \[\mathcal{L} (\theta, \mathcal{D}_P) = \frac{1}{2} \sum_{\mu=1}^P \left[y^\mu - f_{\theta}(x^\mu)\right]^2 \,. \label{eq:loss}\tag{3}\] As a standard practice in statistical mechanics of deep learning, we define the partition function of the problem as: \[Z_{\mathcal{D}_P} = \int d p(\theta) \,e^{-\beta \mathcal{L} (\theta, \mathcal{D}_P)}\;, \label{defpartition}\tag{4}\] where the symbol \(d p(\theta)\) indicates the collective integration of the weights of the network over a prior measure and it will be used interchangeably with \(d\theta\rho(\theta)\), being \(\rho(\theta)\) the prior probability distribution. This has a natural Bayesian learning interpretation: the Gibbs probability \(P_{\beta} (\theta \mid \mathcal{D}_P )=Z^{-1} e^{-\beta \mathcal{L} (\theta, \mathcal{D}_P)} \rho(\theta)\) associated with the partition function in Eq. 4 is the posterior distribution of the weights. To keep the notation concise, we will often omit the explicit dependence on the dataset in the rest of the manuscript. Modeling the standard initialization scheme, the prior distribution over weights is Gaussian: \[W^{(\ell)}_{i_{\ell} i_{\ell-1}} \overset{\mathrm{i.i.d.}}{\sim} \mathcal{N} (0, 1/\lambda)\,, \label{eq:weight95prior}\tag{5}\] where the precision \(\lambda\), the inverse of the variance, is the same at each layer for simplicity of the notation and as usually in practice. In this framework, the average test error over a new (unseen) example \((x^0, y^0)\) is given by: \[\mathbb{E}_{\theta\mid\mathcal{D}_P}[\epsilon_{\textrm g} (x^0, y^0)] = \int d p(\theta) \, [y^0- f_{\theta}(x^0)]^2 \frac{e^{-\beta \mathcal{L}(\theta)}}{Z}\,. \label{eq:err95g95def}\tag{6}\] In numerical experiments, we also consider the empirical generalization error (or simply the generalization error), which is defined as the average test error over \(P_t\) different test examples, where \(P_t\) denotes the number of patterns in the test set. The average training error at a given inverse temperature \(\beta\) is given by: \[\mathbb{E}_{\theta\mid\mathcal{D}_P}[\epsilon_{\textrm t}] = \frac{1}{P} \sum_{\mu=1}^{P}\int d p(\theta) \, \left[y^\mu - f_{\theta}(x^\mu)\right]^2 \frac{e^{-\beta \mathcal{L}(\theta)} }{Z}\,. \label{eq:err95t95def}\tag{7}\] Training and test errors represent two relevant observables that can be computed in the usual way by taking derivatives of the partition function with respect to appropriate source fields.

3 Preamble: the Wishart distribution and its properties↩︎

In statistics, the Wishart distribution arises as a matrix ensemble generalizing the Gamma distribution, and it is defined over symmetric, positive definite random matrices. Let us suppose that \(G\) is a \(P \times N\) matrix, and each column \(G_i\) (\(i = 1, \dots,\, N\)) of \(G\) is independently drawn from a \(P\)-dimensional normal distribution with zero mean and covariance matrix \(V\): \[\begin{align} G&\coloneq(G_1, \dots, G_N), & G_i\overset{\mathrm{i.i.d.}}{\sim} \mathcal{N}_P(0, V). \end{align}\] The Wishart distribution \(\mathcal{W}_P(V,N)\) is the probability distribution of the \(P \times P\) random matrix \[R \coloneq G G^{\top} = \sum_{i=1}^N G_iG_i^\top \sim \mathcal{W}_P(V,N)\,. \label{eq:defWishartEnsemble}\tag{8}\] Using standard terminology in statistics, we will refer to \(N\) as the number of degrees of freedom, and to \(V\) as the scale matrix of the corresponding Wishart distribution. The probability density function of the Wishart distribution is: \[\rho_{\mathcal{W}_P} (R;\, V,N) = \frac{\det(R)^{\frac{(N-P-1)}{2}} \,e^{-\frac{1}{2}\mathrm{Tr}\left(V^{-1} R\right)}}{2^{\frac{NP}{2}} \det(V)^{\frac{N}{2}}\Gamma_P(N/2)}\, ,\] where \(\Gamma_P\) is the multivariate Gamma function. In the special case \(P = V = 1\), note that the Wishart distribution reduces to a chi-squared distribution \(\chi_N^2\) with probability density: \[\rho_{\chi^2}(R;\,N) = \rho_{\mathcal{W}_1}(R;\, 1,N) = \frac{R^{\frac{N}{2}-1} \,e^{-\frac{R}{2}}}{2^{\frac{N}{2}}\, \Gamma (N/2)}\,.\] The mean and variance of a Wishart random matrix \(R \sim \mathcal{W}_P(V/N,\,N)\) are given by \[\begin{align} \label{eq:prop95of95wishart} \mathbb{E} (R_{\mu \nu}) &= V_{\mu\nu}\,,\\ \mathrm {Var}(R_{\mu \nu}) &= \frac{V^2_{\mu\nu} \,+\, V_{\mu\mu}V_{\nu\nu} }{N}\,. \end{align}\tag{9}\] We now highlight an important property of Wishart matrices, which will be fundamental throughout this manuscript: If \(R \sim \mathcal{W}_P(V,\,N)\) and \(C\) is any fixed \(S \times P\) rectangular matrix with rank \(S\), then the \(S\times S\) random matrix \(C R C^\top\) is also Wishart distributed: \[C R C^\top \sim \mathcal{W}_S (C VC^\top,\, N)\,.\] As a corollary, for \(S = 1\) we get that, given a constant \(P\)-dimensional vector \(s\), the scalar random variable \(s^\top R s / s^\top V s\) is chi-squared distributed: \[\frac{s^\top R s}{s^\top V s} \sim \chi_N^2\,. \label{eq:propertychi}\tag{10}\] For properties of the non-central Wishart distribution, see Appendix 9.1.

4 The Equivalent Wishart Ansatz for fully-connected DNNs in the proportional limit↩︎

4.0.1 Prior distribution of network outputs↩︎

As the loss depends on the parameters \(\theta\) only through the network outputs \(f^\mu ,\) and thus the joint output prior \(\rho(f|X)\), the partition function defined in Eq. 4 can be conveniently written in terms of the characteristic function \(\varphi(\bar f|X)= \mathbb{E}_{f}[\exp(-if^\top \bar f)]\) of the output prior: \[\begin{align} Z &= \int \prod_{\mu = 1}^P \frac{df^\mu d\bar f^\mu}{2\pi} e^{-\frac{\beta}{2} \sum_{\mu}\left(y^\mu -f^\mu\right)^2 + i \sum_\mu f^\mu \bar f^\mu} \varphi (\bar f| X)\nonumber \\ &= \langle \varphi(\bar{f}|X) e^{i\sum_\mu y^\mu \bar{f}^\mu} \rangle_{\bar f \sim \mathcal{N}_P(0,\beta \mathbb{1})}, \label{eq:Z95from95priorcharfunc} \end{align}\tag{11}\] where substituting \(f^\mu\to f^\mu + y^\mu\) we performed the Gaussian \(df\) integral, and use \(\langle\cdot \rangle_{\bar f}\) only for brevity of notation of the dual variable measure. We thus need to compute the characteristic function of the joint output prior \(\varphi (\bar f| X)\) by performing the integral over the parameter space: \[\varphi (\bar f| X) = \int d p(\theta) \,e^{-i \sum_{\mu} \bar f^\mu f_\theta(x^\mu)}\,. \label{characteristicoutput}\tag{12}\] Note that the whole complexity of the learning problem lies in computing the prior or its characteristic function - specifically if we can find an integral representation \(\varphi(\bar{f}|X) = \int dQ e^{-S_\varphi(Q,\bar{f})}\) where \(S_\varphi(Q,\bar{f})\) is at most a quadratic function in \(\bar{f}\), and \(Q\) are a set of low-dimensional order parameters, the \(d\bar{f}\) integration to obtain the posterior partition function in Eq. 11 reduces to Gaussian. The EWA leads to just such an integral representation of \(\varphi(\bar{f}|X)\), as shown next.

4.0.2 Output prior as an ensemble of random kernel matrices↩︎

Due to the i.i.d. normal prior of the weights Eq. 5 , the activations Eq. 1 at layer \(\ell\) in the Bayesian network prior are always a sum of Gaussian variables when conditioning on the pre-activations in the layer below. This is to make the standard observation that by introducing the empirical kernels at each layer \(\ell = 1,\,\dots,\, L\): \[K^{(\ell)}_{\mathrm{E}} = \frac{\sigma \left(H^{(\ell)}\right) \sigma \left(H^{(\ell)}\right)^\top}{\lambda N_\ell} \,, \label{empiricalkernel}\tag{13}\] where the activation function \(\sigma\) is applied element-wise to the \(P \times N_\ell\) matrix of pre-activations \(H^{(\ell)},\) and the kernel- or Gram-matrix of the data \[K_{\mathrm{E}}^{(0)}=C=\frac{XX^\top}{\lambda N_0}, \label{eq:def95data95Gram95matrix95C}\tag{14}\] the columns \(H^{(\ell)}_{i_\ell} \coloneq (h^{(\ell)}_{1i_{\ell}}, \dots, h^{(\ell)}_{Pi_{\ell}})^\top\in \mathbb{R}^P\) are conditionally i.i.d. normal distributed: \[\label{eq:preactivations-distribution} H^{(\ell)}_{1}, \dots, H^{(\ell)}_{N_\ell} \mid H^{(\ell-1)} \overset{\mathrm{i.i.d.}}{\sim} \mathcal{N}_P(0,K_{\mathrm{E}}^{(\ell-1)}).\tag{15}\] Correspondingly, also the outputs \(f\) are Gaussian when conditioned on \(K^{(L)}_E\). Since the distributions of \(H^{(\ell)}\) each induce a distribution of \(K^{(\ell)}\) in Eq. 13 that again only depends on the realization of the kernel \(K^{(\ell-1)}\) below, this leads to an alternative description of the prior, and its characteristic function, in terms of \(L\) random matrix ensembles [55]: \[\varphi (\bar f| X) = \prod_{\ell = 1}^L\int_{\mathcal{S}^+_P} dK_{\mathrm{E}}^{(\ell)} \rho (K_{\mathrm{E}}^{(\ell)}|K_{\mathrm{E}}^{(\ell-1)}) e^{-\frac{1}{2}\bar f^\top K^{(L)}_E\bar f}\,, \label{eqprioralternative}\tag{16}\] where the integrations are performed over the cone \(\mathcal{S}_P^+\) of semi-positive definite symmetric \(P\times P\) matrices, and the conditional probabilities \(\rho (K_{\mathrm{E}}^{(\ell)}|K_{\mathrm{E}}^{(\ell-1)})\) are defined as: \[\rho (K_{\mathrm{E}}^{(\ell)}|K_{\mathrm{E}}^{(\ell-1)}) = \mathbb{E}\left[ \delta \bigg( K_{\mathrm{E}}^{(\ell)} - \frac{\sigma (H^{(\ell)}) \sigma (H^{(\ell)})^\top}{\lambda N_\ell} \bigg) \right], \label{conditionaldistrib}\tag{17}\] where the expectation is taken with respect to \(H_i^{(\ell)} \mid H^{(\ell-1)} \overset{\mathrm{i.i.d.}}{\sim} \mathcal{N} (0, K_E^{(\ell -1)})\). Coming from a slightly different perspective, in Ref. [51] a joint (matrix-)Gaussian approximation of these layer-wise distributions is done, leading to a high-dimensional adaptive kernel theory with \(L\) coupled \(P\times P\) dimensional equations of state for the layer-wise empirical kernels in the posterior. Here, we seek an approach to obtain an intensive, low-dimensional equation of state. The fundamental observation that will lead us to formulate the Equivalent Wishart Ansatz arises from the analysis of deep linear networks. In this special case, \(\sigma \left(H\right) \sigma \left(H\right)^\top = H H^\top\) and by Eq. 8 the conditional probability in Eq. 17 , is non-asymptotically characterized as a Wishart distribution: \[\rho_{\mathrm{lin}} (K_{\mathrm{E}}^{(\ell)}|K_{\mathrm{E}}^{(\ell-1)}) = \rho_{\mathcal{W}_P} \left(K_{\mathrm{E}}^{(\ell)};\, K_{\mathrm{E}}^{(\ell-1)}/ N_\ell,\,N_\ell\right)\,.\] This result is implicitly at the core of the drastic dimensional reduction that occurs in finite Bayesian deep linear networks, related to the property Eq. 10 of Wishart distributions.

4.0.3 The Equivalent Wishart Ansatz↩︎

In the non-linear case, we cannot characterize non-asymptotically the conditional probability distributions \(\rho (K_{\mathrm{E}}^{(\ell)}|K_{\mathrm{E}}^{(\ell-1)})\) layer by layer. Still we can hope that simplifications arise in the proportional regime, where \(P,\, N_\ell \to \infty\) and the ratios \(\alpha_\ell = P/N_\ell\) are finite. The fixed reference point is that, starting from the last layer, the mean of the distribution must be the Neural Network Gaussian Process (NNGP) kernel \[\mathbb{E} \left[K_E^{(L)} | K_E^{(L-1)}\right] = \Theta\big(K_E^{(L-1)}\big) \, ,\label{eq:empiricalkernel95mean95conditioned95L-1}\tag{18}\] describing the large-width limit of Bayesian DNNs. The matrix elements of the NNGP kernel are given as \[\left[\Theta(K)\right]_{\mu\nu} = \frac{1}{\lambda} \mathbb{E}_h [ \sigma(h^\mu)\sigma(h^\nu) ] \label{eq:def95NNGP95kernel95one-layer}\tag{19}\] with \(h \sim \mathcal{N}\left(0, \Sigma_{2} \right),\) and \(\Sigma_2 = \left(\begin{smallmatrix} K_{\mu\mu} & K_{\mu\nu}\\ K_{\nu\mu} & K_{\nu\nu} \end{smallmatrix} \right) \in \mathcal{S}_2^+.\) The weight precision \(\lambda\) appears since \(\Theta(K^{(\ell)})_{\mu\nu}=\mathrm{Cov}[h_\mu^{\ell+1}h_\nu^{\ell+1}]\).

The mean in Eq. 18 is again a random variable when moving the conditioning down to the next layer, with mean \(\mathbb{E} \left[\Theta\big(K_E^{(L-1)}\big)| K_E^{(L-2)}\right] = \Theta \circ \Theta\big(K_E^{(L-2)}\big)\big),\) such that by iteration \[\mathbb{E} \left[ \Theta^{L-\ell}\big( K_E^{(\ell)} \big) | K_E^{(\ell-1)} \right] = \Theta^{L-(\ell-1)}\left(K_E^{(\ell-1)}\right). \label{eq:empiricalkernel95mean95conditioned95l}\tag{20}\] Here we introduced the \(l\)-layer NNGP kernel function \[\Theta^\ell(K) = \Theta \circ \dots \circ \Theta (K). \label{eq:def95NNGP95kernel95lth-layer}\tag{21}\]

The Equivalent Wishart Ansatz (EWA) amounts to state that in the proportional regime we can perform analytical calculations as if the fluctuations around each of these means are still Wishart. The Ansatz is thus \[\Theta^{L-\ell}\big( K_E^{(\ell)} \big) \,|\, K_E^{(\ell-1)} \overset{\mathrm{EWA}}{\sim} \mathcal{W}_P \left(V^{(\ell)}, N_\ell\right) \, , \label{eq:EWA-statement}\tag{22}\] with scale matrix (see Eq. 9 ): \[V^{(\ell)} = \frac{1}{N_\ell} \Theta^{L-(\ell-1)}\big( K_E^{(\ell-1)}\big). \label{eq:EWA-scalematix}\tag{23}\] For example \(K_E^{(L)}|K_E^{(L-1)}\sim\mathcal{W} \big( \Theta\big(K_E^{(L-1)}\big)/N_L, N_L\big)\) and \(\Theta^{L-1}\big(K_E^{(1)}\big) \sim \mathcal{W} \big( \Theta^L\big(C\big)/N_1, N_1\big)\). Note that we thereby obtained a new closed sequence of random matrix ensembles \[\rho\big(K_E^{(L)} \big) = \prod_{\ell=L}^1 \rho\left(\Theta^{L-\ell}\big( K_E^{(\ell)} \big) \,|\, \Theta^{L-(\ell-1)}\big( K_E^{(\ell-1)} \big) \right).\] The EWA is rooted in the idea that in first approximation the fluctuations in the kernels of a nonlinear network will be similar to those arising in a linear network. In the proportional limit in particular, for the large majority of vectors \(\bar f\) the contraction \(\bar f^\top K \bar f\) of these random matrices behaves as if \(K\) was Wishart. This is demonstrated numerically in Section 4.2. Note that the EWA is not equivalent to a Gaussian approximation of the network’s pre-activations \(h\), which would strictly yield the lazy NNGP result, nor to a Gaussian approximation of the post-activations at each layer or to a Wishart approximation of each empirical kernel. Rather, propagating fluctuations from each of the lower layer weight priors forward through the nonlinearities, induces prior fluctuations of the last layer empirical kernel, which approximately acts as if Wishart distributed in the proportional regime.

4.0.4 Computation of \(\varphi(\bar f |X)\) given the EWA↩︎

A direct consequence of the EWA is that in Eq. 16 we can exploit the contraction property Eq. 10 and form of the scale matrix Eq. 23 to iteratively substitute \[\begin{align} e^{-\frac{1}{2} \bar{f}^\top K_E^{(L)} \bar{f}} &= \exp\bigg[ -\frac{1}{2} \underbrace{ \frac{\bar{f}^\top K_E^{(L)} \bar{f}}{\bar{f}^\top V^{(L)} \bar{f}}\; }_{\coloneq Q\sim \chi_{N_L}^2} \bar f^\top V^{(L)} \bar f \bigg] \nonumber \\ &= \exp\left[ -\frac{1}{2} \frac{Q_L}{N_L} \bar{f}^\top \Theta(K_E^{(L-1)}) \bar{f} \right] \nonumber \\ &= \exp\left[ -\frac{1}{2} \bigg(\prod_{\ell=1}^L \frac{Q_\ell}{N_\ell} \bigg) \bar{f}^\top \Theta^L(C) \bar{f} \right] \, , \label{eq:iterated95chi295intoduction} \end{align}\tag{24}\] where all \(Q_\ell \sim \chi_{N_\ell}^2\) are independent of each other, of \(\bar{f}\), and of the data Gram matrix \(C\). This makes the EWA extremely powerful in reducing the complexity of the intractable MLP prior while keeping the, we argue, main source of fluctuations in the proportional regime. The characteristic function is thus \[\varphi(\bar{f}|X) \overset{\mathrm{EWA}}{=} \int \big(\prod_\ell dQ_\ell\big) e^{- S_\varphi(Q, \bar{f}) + \sum_\ell \log\big(\rho_{\chi^2}(Q_\ell;\,N_\ell)\big)} \label{eq:charfunc95prior95EWA95singleoutput}\tag{25}\] where \(S_\varphi(Q, \bar{f}) = \frac{1}{2} \big(\prod_{\ell=1}^L \frac{Q_\ell}{N_\ell} \big) \bar{f}^\top \Theta^L(C) \bar{f}\) is only quadratic in \(\bar{f}\), as desired to compute the partition function via Eq. 11 .

In the following, we provide extensive direct and indirect evidence of the validity of the EWA in the proportional limit:

• First, in Section 4.1 we show why the EWA also extends to activation functions with non-zero mean \(\mathbb{E}_h [ \sigma(h^\mu)] > 0\), such as the ReLU activation function. While in principle requiring a generalization to non-central Wishart distributions and leading to significant complications, we show that asymptotically in the proportional limit central and non-central EWA nonetheless yield equivalent predictions, also for non-zero mean activation functions.

• In Section 4.2 we provide a numerical verification of the validity of the EWA; we introduce a Large Deviation Principle and compare the rate function for the variables \(q_\ell \coloneq Q_\ell/N_\ell\) to samples from the true MLP prior. We leverage the simplified expression for the characteristic function, Eq. 25 , to compute the posterior asymptotic free energy density in Section 4.3 and provide an interpretation in terms of data-dependent Gaussian processes.

• In Section 5, we extend this setting to DNNs with multiple outputs and introduce the Stacked Equivalent Wishart Ansatz to describe deep networks with convolutional layers.

• Finally, in Section 6, we present a systematic numerical study comparing the posterior predictor of the EWA on real data against numerical experiments with DNNs.

4.1 EWA in the case of general activation functions.↩︎

In principle, the fact that the mean activation \(m(K) = \mathbb{E}_{h\sim \mathcal{N}(0, K)}[\sigma(h)] \neq 0\) for activation functions such as ReLU suggests that \(\sigma(h)\) should not be replaced by an effective zero-mean variable as supposed by the EWA above. Rather, the natural extension would be to approximate all of the \(\Theta^{L-\ell}(K_E^{(\ell)}) | K_E^{(\ell-1)}\) as noncentral Wishart random matrices. While this ensemble has a more complicated density function, one can iteratively introduce the \(Q_\ell\) variables analogously to Eq. 24 , only that now the \(Q_\ell\) are no longer \(\chi^2\)-distributed. Instead, now the quantity \(\tilde{Q}_{\ell}\coloneq(1 + \lambda_{\mathrm{nc}})Q_{\ell}\) follows a noncentral \(\chi^2\)-distribution \(\tilde{Q}_{\ell} \sim \chi^2_{\mathrm{nc}}(N_\ell, N_\ell \lambda_\mathrm{nc})\) with so-called noncentrality parameter \(N_\ell \lambda_\mathrm{nc}\), where \[\lambda_\mathrm{nc}(\bar{f}, K^{({\ell-1})}_E) = \frac{(\bar{f}^\top m)^2}{\bar{f}^\top \Sigma \bar{f}}.\] See Appendix 9.1 for the derivation. Here \(\Sigma(K) = \mathrm{Cov}[\sigma(h)]_{h\sim \mathcal{N}(0,K)}\) is the covariance kernel instead of the second moment, and for brevity we drop the \(\ell\)-dependent arguments, which play no important role in the following, writing \(\Sigma,m\) and \(\Theta\) instead of \(\Sigma\big(\Theta^{L-\ell-1}(K^{\ell-1}_E)\big)\), \(m\big(\Theta^{L-\ell-1}(K^{\ell-1}_E)\big)\) and \(\Theta^{L-\ell}\big(K^{(\ell-1)}_E\big)\).

Alas, unlike for the central EWA, this distribution still depends on \(\bar{f}\) and \(K^{({\ell-1})}_E\), requiring to introduce an additional order parameter and a complex dependence between the \(Q_{\ell}\) across layers. Indeed the distributions of \(Q_\ell\) and \(Q_{\ell,\mathrm{nc}}\) differ significantly, for example in their variances by a factor \(1 - \frac{\mathrm{Var}[Q_{\ell,\mathrm{nc}}]}{\mathrm{Var}[Q_\ell]} = \frac{(\bar{f}^\top m)^4}{(\bar{f}^\top \Theta\bar{f})^2}\) which is \(O(1)\) even for random overlaps \(\bar{f}^\top m\).

In the lazy infinite-width limit \(N_\ell \to \infty, P=\mathrm{const.}\), these differences vanish since both distributions have the same mean and concentrate in prior and posterior to \(\lim_{N_\ell \to \infty}(Q^\ast_{\ell,nc}/N_\ell) = \lim_{N_\ell \to \infty}(Q^\ast_\ell/N_\ell) = 1\), giving the usual NNGP result. The nonzero mean enters trivially in the form of the NNGP kernel, defined as the covariance of the pre-activations \(h\) and therefore given through the second moment of the activations, \(\Theta = \lambda^{-1}\mathbb{E}[\sigma \sigma^\top] = \lambda^{-1}(\Sigma + m m^\top)\) with \(\Sigma\) the activation covariance, and has no further effects.

In the proportional limit, the fluctuations of \(Q\) become important in the posterior. However, another mechanism suppresses the difference due to the central and noncentral distributions of \(Q\) on the level of both prior and posterior outputs. Intuitively, the reason is that the kernel \(\Theta\) has an \(O(P)\) outlier eigenvalue close to the mean direction caused by the rank-1 spike \(m m^\top\) which suppresses any contributions to the prior from \((\bar{f}^\top m)^2 > O(1)\) under the \(d\bar f\) integral. The denominator in contrast concentrates to \(\bar{f}^\top \Sigma \bar{f} \sim P\), such that only noncentrality parameters \(\lambda_\mathrm{nc}(\bar{f}, K^{({\ell-1})}_E)= O(1/P)\) contribute to both prior and posterior; any contributions where \(\lambda_\mathrm{nc} \neq 0\) and \(\rho(Q_{\mathrm{nc}})\) would asymptotically differ from \(\chi^2(N_\ell)\) are removed.

In the following we detail an argument showing this asymptotic equivalence holds at any layer, both in the prior- and the posterior output distribution. Consider the contribution to the partition function from any single layer \(l\)\[\begin{align} Z(K^{(\ell-1)}_E) &= \big\langle e^{i y^\top \bar{f}} \varphi(\bar{f}|K^{(\ell-1)}_E) \big\rangle_{\bar f \sim \mathcal{N}(0 , \beta \mathbb{1}_P)} \label{eq:noncentra-central95Z}\\ &= \int d\bar{f}\, d\tilde{Q}_{\ell}\; \rho_{\chi^2_{\mathrm{nc}}}\big( \tilde{Q}_{\ell}; N_\ell, N_\ell \lambda_\mathrm{nc}(\bar{f}, K^{({\ell-1})}_E) \big) \nonumber \\ & \qquad \quad \times e^{- \frac{1}{2} \bar{f}^\top \big(\beta^{-1} \mathbb{1}+ Q_{\ell}\Theta \big) \bar{f} + i y^\top \bar{f}} \, . \nonumber \end{align}\tag{26}\] Note that Eq. 26 is of the form \[\int d\tilde{Q} \left\langle F_{\tilde{Q}}\left( \lambda_\mathrm{nc}(\bar f) \right) \right\rangle_{\bar f \sim \mathcal{N}(\mu, K_\beta)} \label{eq:noncentral95function95under95fbar95measure}\tag{27}\] with \(\lambda_\mathrm{nc}(\bar f) =\frac{(\bar{f}^\top m)^2}{\bar{f}^\top \Sigma \bar{f}}\), covariance \(K_\beta =\big[ \beta^{-1} \mathbb{1}+ Q(\Sigma + m m^\top) \big]^{-1}\) and \(\mu = i K_\beta\, y\). Here \(F\) depends on \(\bar f\) only through \(\lambda_\mathrm{nc}(\bar f)\) and is a positive continuous function. Assuming that \(\Sigma\) has an extensive number of eigenvalues above the noise floor \(\beta^{-1}\), we now show that \(\lambda_{nc}\) is self-averaging and vanishes as \(O(1/P)\) under the \(\bar f\) measure, such that \[\left\langle F_{\tilde{Q}}\left( \lambda_\mathrm{nc}(\bar f) \right) \right\rangle_{\bar{f}} = F_{\tilde{Q}}\left( \langle\lambda_\mathrm{nc}(\bar f) \rangle_{\bar f} \right) = F_{\tilde{Q}}\left( O(1/P) \right),\] and we can replace \(d\tilde{Q}_{\ell} \rho_{\chi^2_{\mathrm{nc}}}\big( \tilde{Q}_{\ell}\big) \to dQ_\ell\, \rho_{\chi^2}(Q_\ell)\) in Eq. 26 , recovering Eq. 25 .

Given that the numerator \((\bar f^\top m)^2\) only depends on a single projection \(m\) and this direction contributes negligibly to \(\bar f^\top \Sigma \bar f\), self-averaging of the numerator and denominator can be assessed separately. In the numerator, \((\bar f^\top m)^2 =O(1)\) for typical \(\bar f \sim \mathcal{N}(\mu, K_\beta)\) due to the small inverse outlier eigenvalue \(\sim 1/\| m\|^{2}\) asymptotically localized to the \(m\) direction. A convenient scalar controlling the typical size of the denominator is the effective dimension \[t \coloneq \mathop{}\!\mathrm{Tr}\left[\Sigma \, K_\beta \right] = \mathop{}\!\mathrm{Tr}\left[ \Sigma \left(\beta^{-1}\mathbb{1}+ Q(\Sigma + m m^\top) \right)^{-1}\right]. \label{eq:teff95def}\tag{28}\] A short calculation in Appendix 9.4 shows that \(\bar f^\top \Sigma \bar f = O(t)\) with relative fluctuations \(\sim t^{\frac{1}{2}}\), and is therefore self-averaging if \(t = O(P)\). Indeed, in Eq. 28 at low temperature \(\beta^{-1}\ll1\) and aside from the negligible \(m\) direction, \(\Sigma\) and \(K_\beta\) admit almost the same diagonalization. Along an eigenmode with \(\Sigma\)-eigenvalue \(\lambda_i\) and corresponding \(K_\beta\)-eigenvalue \([\beta^{-1}+Q\lambda_i]^{-1}\), the contribution to \(t\) is then \(\frac{\lambda_i}{\beta^{-1}+Q\lambda_i} z_i^2\) with \(z_i\) standard normal. Hence for the extensive number of \(\lambda_i \gtrsim \beta^{-1}\) modes one has \(\lambda_i(\beta^{-1}+Q\lambda_i)^{-1}=O(1)\) and their contributions add up to \(t \sim P,\) with relative fluctuations of order \(P^{-1/2}\). In particular, the denominator \(\bar f^\top \Sigma \bar f\) then concentrates away from zero, so \(1/(\bar{f}^\top \Sigma \bar f)\) is also self-averaging and \(\langle1/(\bar{f}^\top \Sigma \bar f)\rangle_{\bar f} \sim 1/t=O(1/P)\). Combining numerator and denominator we have for the dominant \(\bar f\) configurations \(\lambda_\mathrm{nc} = O(1/P)\) and correspondingly \[\rho_{\chi^2_{\mathrm{nc}}}\!\Big(\tilde{Q}_{\ell};N_\ell, N_\ell\lambda_{\mathrm{nc}}\Big) = \rho_{\chi^2}\!\big(Q_{\ell};N_\ell\big)\Big(1+O(1/P)\Big).\]

These arguments can be applied layer-by-layer, first replacing \(\rho_{\chi^2_{\mathrm{nc}}}\big( \tilde{Q}_{L}\big) \to \rho_{\chi^2}(Q_L)\) for \(\ell = L\), then introducing \(\tilde{Q}_{L-1}\) as in Eq. 24 and repeating the replacement, down to \(\ell=1\). This asymptotic equivalence of noncentral and central EWA not only holds at the level of the posterior partition function Eq. 11 , but also for the prior output distribution \(\rho(f|X)\), as can be seen by setting \(\beta\to\infty\) in Eq. 26 such that the computation corresponds to the Fourier back-transform of the prior characteristic function.

The more complicated noncentral EWA therefore reduces to the zero-mean Ansatz, and is not considered further. Only at \(P \approx 50\) we could find small differences between noncentral and zero-mean theories as a finite-size effect (not shown). For additional numerical evidence and discussion of the mean contribution, see Appendix 9.

4.2 Numerical validation of the EWA↩︎

The aim of this section is to provide numerical evidence for the EWA using the framework of Large Deviation Theory. As discussed in the previous section, possible contributions due to a non-zero mean of the activation function are negligible in the asymptotic regime. Therefore, it is sufficient to investigate whether a Large Deviation Principle (LDP) holds for the central EWA, which states that: \[\Theta^{L-\ell}(K_\mathrm{E}^{(\ell)}) \mid K_\mathrm{E}^{(\ell-1)} \sim \mathcal{W}_P \left( \frac{\Theta^{L-\ell+1}(K_\mathrm{E}^{(\ell-1)})}{N_\ell}, N_\ell \right),\] where \(\Theta\) is the usual NNGP kernel function. Assuming the validity of the EWA implies that property in Eq. 10 needs to be verified, i.e. that, at each layer, the variables \(Q_\ell\) defined as \[Q_\ell = N_\ell \frac{\bar f^\top \Theta^{L-\ell}(K_\mathrm{E}^{(\ell)}) \bar f}{\bar f^\top \Theta^{L-\ell+1}(K_\mathrm{E}^{(\ell-1)}) \bar f}\] are chi-squared distributed for every fixed realization of the \(\bar f \in \mathbb{R}^P\) (conditioned on the empirical kernel at the previous layer): \[Q_\ell \mid K_\mathrm{E}^{(\ell-1)} \sim \chi^2_{N_\ell} = \Gamma(N_\ell/2, 2).\] Notice that the distribution of \(Q_\ell \mid K_\mathrm{E}^{(\ell-1)}\) actually does not depend on \(K_\mathrm{E}^{(\ell-1)}\), implying that the \(Q_\ell\) variables are decoupled across layers, and it is also independent on \(\bar{f}\). Since \(\mathbb{E}[Q_\ell] = N_\ell\), let us define the corresponding intensive variables \(q_\ell = Q_\ell / N_\ell \sim \Gamma(N_\ell/2, 2/N_\ell)\), which concentrates around \(q_\ell=1\) in the infinite-width limit. The intensive variables \(q_\ell\), each viewed as a sequence over the corresponding \(N_\ell\), satisfy a LDP of the form: \[\begin{align} \rho_{q_\ell}(x) &\sim e^{-a_{N_\ell} \mathcal{I}_\ell(x)} & &\text{as } N_\ell \to \infty. \end{align}\] \(\mathcal{I}_\ell(x)\) is called rate function, and it is calculated via the Gärtner-Ellis theorem as the Legendre transform of the scaled cumulant generating function \(\Lambda_{q_\ell}\): \[\begin{align} \mathcal{I}_{q_\ell}(x) &= \sup_t \{ tx - \Lambda_{q_\ell}(t) \}, \\ \Lambda_{q_\ell}(t) &\coloneq \lim_{N_\ell \to \infty} \frac{1}{a_{N_\ell}} \ln M_{q_\ell}(a_{N_\ell} t). \end{align}\] Using for \(M_{q_\ell}\) the moment generating function of a \(\Gamma(N_\ell/2, 2/N_\ell)\) distribution, it follows that the scales for these LDPs are \(a_{N_\ell}=N_\ell\) and that the rate function is the same for all layers: \[\label{eq:single-q-rate-func} \mathcal{I}_{q_\ell}(x) = \mathcal{I}_q(x) = \frac{1}{2} (x-1-\ln x).\tag{29}\] If we take \(N_\ell=N \, \forall l\), the random variables \(q_\ell\) are all independent and identically distributed, and we can derive a multivariate LDP with a single scale \(a_N=N\) (otherwise a multi-scale LDP should be considered): \[\begin{align} \label{eq:LDP-multivariate-q} \rho_{(q_1,\dots,q_L)}(x_1, \dots, x_L) &\sim e^{-N \sum_{\ell=1}^L \mathcal{I}_q(x_\ell)} & &\text{as }N\to\infty. \end{align}\tag{30}\] As shown in the previous section, see in particular Eq. 25 , the characteristic function of the prior over the network output actually depends only on the product \(\mathcal{Q}\coloneq\prod_{\ell=1}^L q_\ell\). We can derive the rate function for the random variable \(\mathcal{Q}\) using the contraction principle: \[\begin{align} \label{eq:product-rate-function} \mathcal{I}_{\mathcal{Q}}(y) &= \inf \left\{ \sum_{\ell=1}^L \mathcal{I}_q(x_\ell) : \prod_{\ell=1}^L x_\ell = y \right\} \\ &= L \, \mathcal{I}_q(y^{1/L}). \end{align}\tag{31}\]

Fig. 2 shows the rate function for the \(\mathcal{Q}\) variable for both Erf and ReLU activation functions, and for several networks of different depth. Black lines denote the theoretical rate function obtained under the EWA, Eq. 31 , while the colored points correspond to actual samples of the \(\mathcal{Q}\) variables, in particular for the CIFAR-10 dataset. The plots suggest in all cases convergence to the theoretical-asymptotic result, even for deeper networks, by showing the empirical rate function for increasing values of number of data points \(P\) and widths of the networks, with constant ratio \(\alpha=1.0\). It is important to underline that an independent sampling of the \(q_\ell\) variables is required in order to build the proper \(\mathcal{Q}\) variable satisfying the LDP with rate function in Eq. 31 , in addition to keeping the \(\bar f\) vector fixed across samples. About the sampling procedure, the independent samples of the vectors \((q_1,\dots,q_L)\) are obtained by sampling Gaussian preactivations, according to Eq. 15 , and building the corresponding empirical kernels at each layer. The NNGP kernel function \(\Theta\) is then applied recursively to build the \(q_\ell\) variables at each layer, and this entire procedure is repeated for each sample keeping the same \(\bar f\). At the end, the \(\mathcal{Q}=\prod_\ell q_\ell\) variables (one for each network) are calculated from the independent samples. In Fig. 3 we show the robustness of the EWA across the different datasets and for different values of the \(\alpha\) parameter. The conclusions are the same as for Fig. 2 both for the MNIST dataset (first column) and for Gaussian data (second columns), as well as for \(\alpha<1\), \(\alpha=1\) and \(\alpha>1\). For more details about the large deviation analysis, a layer by layer analysis of the LDP, and more details about the non central case, we refer the reader to App. 10.

4.3 Effective action for MLPs of depth \(L\)↩︎

Using the EWA expression for the prior characteristic function, Eq. 25 , a Gaussian integration gives the partition function Eq. 11 as the following expectation value over the \(\{q_\ell\}_\ell\) variables: \[\label{eq:Z-as-expectation-for-Varadhan} Z = \mathbb{E}_{q_1,\dots,q_L} \left[e^{-\frac{N}{2} U(q_1,\dots,q_L)}\right],\tag{32}\] where \(U(q) \coloneq U(q_1, \dots, q_L)\) is given by: \[\begin{align} U(q) &= \frac{\alpha}{P} \log \det \left[\mathbb{1}+ \beta K^{(\text{R})}_{\mathcal{Q}}\right] + \\ & \quad \quad \quad + \frac{\alpha}{P} y^\top \left[\beta^{-1}\mathbb{1}+ K^{(\text{R})}_{\mathcal{Q}}\right]^{-1}y, \end{align}\] and the renormalized kernel is: \[\begin{align} \label{eq:renormalized95kernel} K^{(\text{R})}_{\mathcal{Q}} &= \mathcal{Q} \Theta^L(C), & \mathcal{Q}&=\prod_{\ell=1}^L q_\ell. \end{align}\tag{33}\]

Since the \(\{q_\ell\}_\ell\) satisfy the joint LDP given by Eq. 30 , we can use Varadhan’s Lemma (see for example Ref. [56]) to directly calculate the asymptotic free energy density of the system: \[\label{eq:min95of95effective95action} - \frac{1}{N} \ln Z \to \inf_{q_1, \dots, q_L} \underbrace{\left[\sum_{\ell=1}^L \mathcal{I}_q(q_\ell) + U(q_1, \dots, q_\ell)\right]}_{S(q)}.\tag{34}\] We notice that, in order to make this step we need to make the technical assumption that the scaling of the two terms in \(U\) are well behaved in the proportional limit, when \(P\to\infty\), and in particular that the limit exists and it is finite. We also point out that the application of Varadhan’s Lemma and of LDT in general is a more powerful method than the calculation of \(Z\) trough saddle point/Laplace approximation, which would formally yield the same result, because it does not make assumptions of the convergence properties of the partition function itself, but directly provide the asymptotic free energy density, which is, after all, the only quantity we are actually interested in. Indeed, there exist sequences of random variables for which no density–or no useful asymptotic density–is available, yet they satisfy a large deviation principle at the level of measures, and this is sufficient to Varadhan’s lemma to be applicable. In this context, it is indeed very hard to make mathematically rigorous statements about the partition function itself in the proportional limit. Note that in Eq. 34 we introduce a new quantity \(S(q)\), called effective action in the physics language, whose minimum is the free energy density of the system and which is given by (up to irrelevant additive constants): \[\begin{align} \label{eq:action95singleoutput} S(q) = \sum_{\ell=1}^L [q_\ell - \log q_\ell] & + \frac{\alpha}{P} \log \det \left[\mathbb{1}+ \beta K^{(\text{R})}_{\mathcal{Q}}\right] + \\ & + \frac{\alpha}{P} y^\top \left[\beta^{-1}\mathbb{1}+ K^{(\text{R})}_{\mathcal{Q}}\right]^{-1}y. \end{align}\tag{35}\] We notice that the obtained free energy, corresponding to the marginal log-likelihood in Bayesian language, is that of a hierarchical Gaussian process regression model, where the kernel is being renormalized by a function of the order parameters \(\{q_\ell\}_\ell\), which are in principle distributed according to their own hyper-prior but that concentrates around some specific data-dependent values in the proportional limit (for a detailed discussion in the 1HL case, we refer the reader to Ref. [57]). Indeed, in the effective action, the first term corresponds to the order-parameter hyper-prior (and physically to the entropic contribution in the free energy), while the function \(U(q_1, \dots, q_\ell)\) is exactly the marginal log-likelihood of a Gaussian process regression model, conditioned on the value of the hyperparameters \(\{q_\ell\}_\ell\). This implies the possibility to do exact inference with this theoretical model, and in particular to find a closed form expression for the posterior predictive distribution and therefore for relevant observables like the generalization error, see for example Ref. [58]. As discussed in Sec. 4.2, the \(\{q_\ell\}_\ell\) variables have expectation value \(\mathbb{E}[q_\ell]=1\) and variance \(\mathrm{Var}[q_\ell]=N_\ell^{-1}\), such that they concentrate in the proportional limit. The action Eq. 35 governing the posterior however, is minimized in general by values of \(q^\ast_\ell \ne 1\), thus far in the tail of the prior distribution. This arises from the fact that the likelihood (loss term) scales as \(\sim P\), adding a constraint for each data sample. The dominant contribution which a good approximation of the network prior must capture is therefore not the mode of the distribution, but its large deviations behavior.

5 Kernel Renormalization schemes for deep architectures with Multiple Outputs and Convolutional Layers↩︎

5.1 Fully-connected DNNs with multiple outputs↩︎

The multiple outputs architecture implements a function \(f_\theta : \mathbb{R}^{N_0} \to \mathbb{R}^D\), with \(D>1\). The last-layer weights thus form a \(D \times N_L\) matrix, whose elements are labeled by an additional index, leading to the following definition of the network function: \[f_{a}(x) = \frac{1}{\sqrt{N_L}} \sum_{i_L=1}^{N_L}W^{(L+1)}_{a i_L} \sigma\Bigl(h^{(L)}_{i_L}\Bigl) \, ,\] where the pre-activations \(h^{(L)}_{i_L}\) are defined in Eq. 1 . The square-loss function is the natural generalization of Eq. 3 , replacing the square with the squared Frobenius norm. Let us start from the 1HL case. The EWA amounts to assume that the characteristic function of the prior over outputs is given by: \[\begin{align} \varphi(\bar{f}|X) &= \int_{S_P^+} dp(K_\mathrm{E}) e^{-\frac{1}{2} \mathop{}\!\mathrm{Tr}(\bar{f} K_\mathrm{E} \bar{f}^\top)}, \\ K_\mathrm{E} \mid C &\sim \mathcal{W}_P(\Theta(C)/N_1, N_1). \end{align}\] In the last equation, \(\bar{f} \in \mathbb{R}^{D \times P}\) denotes the matrix whose elements are the dual variables \(\bar{f}_a^\mu\). To proceed with the calculation, we will need two mathematical identities:

• Ingham-Siegel integral [59], [60], or Laplace transform of a Wishart distribution:
  \[\label{eq:Ingham-Siegel} \int_{S_P^+}dM \, \rho_{\mathcal{W}_P}(M | V,N) e^{-\frac{\alpha}{2} \mathop{}\!\mathrm{Tr}(AM)} = [\det(\mathbb{1}_P + \alpha VA)]^{-N/2} \,.\tag{36}\]

• Weinstein–Aronszajn identity [61]: given \(A\) and \(B\) matrices of size \(m \times n\) and \(n \times m\) respectively, the following identity holds \[\label{eq:Weinstein-Aronszajn-identity} \det (\mathbb{1}_m + A \cdot B) = \det (\mathbb{1}_n + B \cdot A).\tag{37}\]

After using the cyclic property of the trace operation, we can proceed with the calculation as \[\begin{align} \varphi(\bar{f}|X) &= \left[ \det \left( \mathbb{1}_P + \frac{1}{N_1} \Theta(C) \bar{f}^\top \bar{f} \right) \right]^{-N_1/2} \\ &= \left[ \det \left( \mathbb{1}_D + \frac{1}{N_1}\bar{f} \Theta(C) \bar{f}^\top \right) \right]^{-N_1/2} \\ &= \int_{S_D^+} dp(q) e^{-\frac{1}{2} \mathop{}\!\mathrm{Tr}(q \bar{f} \Theta(C) \bar{f}^\top)} \\ &= \int_{S_D^+} dp(q) e^{-\frac{1}{2} \bar{f} (q \otimes \Theta(C)) \bar{f}^\top}, \end{align}\] where \(q\) is a Wishart distributed random matrix, \(q\sim\mathcal{W}_D(\mathbb{1}_D/N_1, N_1)\). To understand how to generalize this result to the deep case, let us focus to a 2HL architecture, for which the EWA reads (recall that \(K_\mathrm{E}^{(0)} = C\)):

\[\label{eq:2HL-EWA-multiple-output} \begin{align} K_E^{(2)} \mid K_E^{(1)} &\sim \mathcal{W}_P\left(\Theta(K_E^{(1)})/N_2, N_2\right), \\ \Theta(K_E^{(1)}) \mid K_E^{(0)} &\sim \mathcal{W}_P\left(\Theta^2(C)/N_1, N_1\right) . \end{align}\tag{38}\] The characteristic function of the prior over outputs is in this case given by: \[\varphi(\bar{f}| X) = \int dp(H^{(1)}) \int_{S_P^+} dp(K_\mathrm{E}^{(2)} \mid H^{(1)}) \, e^{-\frac{1}{2} \mathop{}\!\mathrm{Tr}(\bar{f} K_\mathrm{E}^{(2)} \bar{f}^\top)},\] where in the distribution of \(K_\mathrm{E}^{(2)} \mid H^{(1)}\) the dependence on the pre-activations \(H^{(1)}\) is only through the empirical kernel \(K_\mathrm{E}^{(1)}\), allowing for a simple change of the integration variable to \(K_\mathrm{E}^{(1)}\). The inner integral has the same structure as in the 1HL case, leading to: \[\begin{align} \varphi(\bar{f}|X) = \int_{S_P^+} d p(K_\mathrm{E}^{(1)}) \int_{S_D^+} dp(q_2) e^{ -\frac{1}{2} \mathop{}\!\mathrm{Tr}(q_2 \bar{f} \Theta(K_E^{(1)}) \bar{f}^\top) } \, , \end{align}\] where \(q_2 \sim \mathcal{W}_D\left(\mathbb{1}_D / N_2, N_2 \right)\). Given that \(q_2\) is a Wishart matrix, it is symmetric, positive-definite. Therefore we can decompose it (using for example the Cholesky decomposition) as \(q_2=U_2U_2^\top\). In this way, we can write the exponent as \[\mathop{}\!\mathrm{Tr}(q_2 \bar{f} \Theta(K_E^{(1)}) \bar{f}^\top) = \mathop{}\!\mathrm{Tr}( \underbrace{U_2^\top \bar{f}}_{g} \Theta(K_E^{(1)}) \underbrace{\bar{f}^\top U_2}_{g^\top}) \, .\] At this point, we can apply directly the 2HL EWA, Eq. 38 , obtaining the same integral expression as in the last-layer (up to a replacement \(\bar f \to g\coloneq U_2^\top \bar f\)), leading to: \[\begin{align} \varphi(\bar{f}|X) & = \int_{(S_D^+)^2} dp(q_1, q_2) e^{ -\frac{1}{2} \bar{f} (\mathcal{Q} \otimes \Theta^2(C)) \bar{f}^T}, \\ \mathcal{Q} &\coloneq U_2U_1U_1^TU_2^T, \\ q_\ell & \sim \mathcal{W}_D\left(\mathbb{1}_D / N_\ell, N_\ell \right), \quad U_\ell U_\ell^T = q_\ell. \end{align}\] It is now clear that the same procedure can be carried on also for deeper networks, leading to the following kernel renormalization scheme: \[\begin{align} K^{(\text{R})}_{\mathcal{Q}} &= \mathcal{Q} \otimes \Theta^L(C), \\ \mathcal{Q} &= \left( \prod_{\ell=1}^{L}U_{\ell}^\top \right)^\top\left( \prod_{\ell=1}^{L}U_{\ell}^\top \right) \, . \end{align}\] Notice that, in the linear case, we recover the same result of Refs. [23], [24].

5.2 The stacked Equivalent Wishart Ansatz for DNNs with convolutional layers↩︎

One additional advantage of the EWA formalism is that it can be used to obtain, for the first time, the kernel renormalization scheme for non-linear deep neural networks with convolutional layers (CNNs). Here we analyze the general case for arbitrary stride and filter sizes. Also the input data have the usual spatial dimension and display different channels, i.e. each input has the form \(x^\mu_{i_0 a_0}\), where \(a_0 = 1, \ldots, A_0\) (with \(A_0\) being the total number of input channels) and \(i_0 = 1, \ldots, N_0\). In the first layer, convolutional pre-activations are defined as \[h_{i_{1} a_{1}}^{(1)^\mu}=\frac{1}{\sqrt{M A_{0}}}\sum_{a_{0}=1}^{A_{0}}\sum_{m=-\lfloor M/2 \rfloor }^{\lfloor M/2 \rfloor}W_{a_{1}a_{0}m}^{(1)}x_{S_1 i_{1}+m, a_{0}}^{\mu}\,,\] where \(M\) is the dimension of the mask, \(S_1\) is the stride at the first layer, and \(a_0\) and \(a_1 = 1, \ldots, A_1\) denote the input and output channel indices, respectively. The index \(i_1\) runs over the number of patches in the first layer, i.e., \(i_1 = 1, \ldots, N_{p_1} = \lfloor N_0 / S_1 \rfloor\). In the other layers, given the strides \(S_\ell\), each internal kernel has its own patch index \(i_\ell\), where \(i_\ell = 1, \ldots, N_{p_\ell} = \lfloor N_{p_{\ell-1}} / S_\ell \rfloor\). At each layer \(\ell>1\), pre-activations are defined as: \[h_{i_{\ell} a_{\ell}}^{(\ell)^\mu}=\frac{\sum_{a_{\ell-1}=1}^{A_{\ell-1}}\sum_{m=-\lfloor M/2 \rfloor }^{\lfloor M/2 \rfloor}W_{a_{\ell}a_{\ell-1}m}^{(\ell)}\sigma\Bigl(h^{(\ell-1)^\mu}_{S_\ell i_{\ell}+m, a_{\ell-1}}\Bigl)}{\sqrt{M A_{\ell-1}}}\,,\] where \(A_{\ell-1}\) is the number of input channels at layer \(\ell-1\), and \(i_\ell = 1, \ldots, A_\ell\) is the output channel index of the \(\ell\)-th layer. Let us first consider 1HL CNNs; such architectures implement the following function: \[f_{\theta}(x^{\mu})=\frac{{1}}{\sqrt{{A_{1}N_{p_1}}}}\sum_{a_{1}=1}^{A_{1}}\sum_{i_{1}=1}^{N_{p_1}}W_{a_{1}i_{1}}^{(2)}\sigma\left(h_{i_{1}a_{1}}^{(1)^{\mu}}\right)\, .\] Here it is useful to introduce the stacked pre-activation matrix (note that this is a generalization of the fully-connected case): \[\begin{align} \label{eq:stacked-H} H^{(1)} & = \left(\begin{array}{c} h_{1:N_{p_1},1}^{(1)^{1:P}}, \ldots , h_{1:N_{p_1},A_1}^{(1)^{1:P}} \end{array}\right) \nonumber \\ & = \left(\begin{array}{c} H^{(1)}_1, \ldots , H^{(1)}_{A_1} \end{array}\right) \in\mathbb{R}^{N_{p_1}P\times A_{1}}\, , \end{align}\tag{39}\] where \(h_{1:N_p,a_1}^{(1)^{1:P}}\) are \(\mathbb{R}^{N_{p_1}P}\)-dimensional vectors obtained by flattening first the \(\mu\) indices and then the \(i_1\) indices. It is important to note that the stacked matrix contains \(A_1\) independent Gaussian \(\mathbb{R}^{N_{p_1}P}\)-dimensional vectors, with covariance matrix \(G \in \mathbb{R}^{N_{p_1}P \times N_{p_1}P}\) defined by \[\begin{align} G & = \left(\begin{array}{cccc} G_{11} & G_{12} & \dots & G_{1N_{p_1}}\\ G_{21} & \ddots & & \vdots\\ \vdots & & \ddots & \vdots\\ G_{N_{p_1}1} & \ldots & \ldots & G_{N_{p_1}N_{p_1}} \end{array}\right)\, , \\ [G_{i_1 i_1'}]^{\mu\nu} & = \frac{{1}}{MC_{0}}\sum_{a_{0}m}x_{a_{0},S_1i_{1}+m}^{\mu}\,x_{a_{0},S_1i_{1}'+m}^{\nu} \, . \label{eq:def-g-cnn-1hl} \end{align}\tag{40}\] After the readout integration, we can recast the characteristic function of the prior using the pre-activations in Eq. 39 in the following way: \[\begin{align} \label{eq:characteristic-function-prior-cnn} & \varphi(\bar{f}|X) = \int\prod_{a_1} dH^{(1)}_{a_1} \rho_\mathcal{{N}}\left(H^{(1)}_{a_1} ; 0,G \right) \\ & \quad e^{-\frac{{1}}{2\lambda N_{p_1}}\sum_{i_{1}}\left[\sum_{\mu\nu}\bar{f}^{\mu}\left(\frac{1}{A_{1}}\sum_{a_{1}}\sigma\left(h_{i_{1}a_{1}}^{(1)^{\mu}}\right)\sigma\left(h_{i_{1}a_{1}}^{(1)^{\nu}}\right)\right)\bar{{f}}^{\nu}\right]}\nonumber\, . \end{align}\tag{41}\] Introducing the stacked matrix in Eq. 39 has two main advantages: (i) it directly leads to the stacked definition of the empirical kernel at the first layer, \[K_E^{(1)} = \frac{1}{A_1 \lambda}\sigma(H^{(1)})\sigma(H^{(1)})^\top \, ,\] in terms of which we can state the stacked Equivalent Wishart Ansatz, i.e., \[\begin{align} K_E^{(1)} \sim \mathcal{W}_{N_{p_1} P}(\Theta(G)/A_1, A_1) \, , \\ A_1,P\to \infty,\;\alpha=P/A_1=\text{const.}\, ; \end{align}\] (ii) It enables a rewriting of the exponential term in Eq. 41 , such that \[\sum_{i_1 \mu \nu}\bar{f}^{\mu}\left(\frac{1}{A_{1}\lambda}\sigma\left(H_{i_{1}}^{(1)}\right)\sigma\left(H_{i_{1}}^{(1)}\right)\right)\bar{f}^{\nu} = \mathop{}\!\mathrm{Tr}\Bigl[\tilde{F}^\top K_E^{(1)} \tilde{F} \Bigr]\, .\] In the equations above, we also introduced the stacked scale matrix as \[\begin{align} \Theta(G) & = \left(\begin{array}{cccc} \Theta(G)_{11} & \Theta(G)_{12} & \dots & \Theta(G)_{1N_{p_1}}\\ \Theta(G)_{21} & \ddots & & \vdots\\ \vdots & & \ddots & \vdots\\ \Theta(G)_{N_{p_1}1} & \ldots & \ldots & \Theta(G)_{N_{p_1}N_{p_1}} \end{array}\right)\, , \\ [\Theta(G)_{ij}]^{\mu\nu} & = \langle \sigma(h^\mu)\sigma(h^\nu)\rangle_{h\sim\mathcal{N}(0,\Sigma_{\text{stack}})}\, , \end{align}\] where \(\Sigma_{\text{stack}}=\left(\begin{smallmatrix} G_{ii}^{\mu\mu} & G_{ij}^{\mu\nu}\\ G_{ji}^{\nu\mu} & G_{jj}^{\mu\mu} \end{smallmatrix} \right)\), and the stacked matrix of dual outputs, \[\tilde{F} = \left(\begin{array}{cccc} \bar{f} & 0 & \dots & 0\\ 0 & \bar{f} & & \vdots\\ \vdots & 0 & \ddots & \vdots\\ 0 & \ldots & \ldots & \bar{f} \end{array}\right)\;\in\mathbb{R}^{N_{p_1}P\times N_{p_1}}\, .\] It is worth mentioning that in the CNN case, the definition of the empirical kernel leads to a new natural definition of the proportional regime, where the quantities that are taken to be large are the number of patterns and the number of channels (this is also in agreement with Ref. [62], where the large-width limit of CNNs was first introduced). Taking advantage of the stacked notation, we can compute the characteristic function of the prior: \[\begin{align} \tag{42} \varphi(\bar{f}|X) & = \int_{\mathcal{S}_{N_{p_1} P}^+} dK_E^{(1)} \rho_{\mathcal{W}_{N_{p_1} P}}(K_E^{(1)}; \Theta(G)/A_1, A_1 ) \nonumber \\ & \quad\quad\quad\quad\quad\quad\quad\quad \times e^{-\frac{1}{2N_{p_1}}\mathop{}\!\mathrm{Tr}\Bigl[\tilde{F}^\top K_E^{(1)} \tilde{F} \Bigr]} \\ & = \det\Bigl[\mathbb{1}_{N_{p_1} P} + \Theta(G) \tilde{F} \tilde{F}^\top / (N_{p_1} A_1) \Bigr]^{-\frac{A_1}{2}} \tag{43} \\ & = \det\Bigl[\mathbb{1}_{N_{p_1}} + \tilde{F}^\top \Theta(G) \tilde{F} / (N_{p_1} A_1) \Bigr]^{-\frac{A_1}{2}}\, , \tag{44} \end{align}\] where in the second line we used the Ingham–Siegel integral in Eq. 36 , and in the last line the Weinstein–Aronszajn identity in Eq. 37 . From the last line, the dimensional reduction can be made explicit using the inverse Laplace transform of a Wishart distribution: \[\label{eq:1hl-prior-cnn} \varphi(\bar{f}|X) = \int_{\mathcal{S}_{N_{p_1}}^+} dp(q) \;e^{-\frac{1}{2N_{p_1}}\mathop{}\!\mathrm{Tr}(\tilde{F}^\top \Theta(G) \tilde{F}q)}\, ,\tag{45}\] where the low-dimensional order parameter \(q \sim {\mathcal{W}_{N_{p_1}}}(\mathbb{1}/ A_1, A_1 ) \in \mathbb{R}^{N_{p_1}\times N_{p_1}}\) naturally emerges. The kernel renormalization scheme is, in this case, \[K^{(\text{R})}_q = \frac{1}{N_{p_1}}\sum_{ij=1}^{N_p}q_{ij}\Theta(G)_{ij} \, .\] Note that in the special case of a single input channel \(A_0 = 1\), the kernel renormalization scheme following from the stacked EWA reproduces the local kernel renormalization found in Ref. [43]. On the contrary, in the large-width limit where \(q\) becomes deterministic around \(q = \mathbb{1}_{N_{p_1}}\), we recover the infinite-width results first introduced in Ref. [62]. It is worth mentioning that since the characteristic function of the prior is still a quadratic form in \(\bar{f}\), it is very easy to compute the corresponding posterior effective action following the guidelines outlined in Sec. 4.3.

To extend our results to deep CNNs, we now consider a 2HL CNN, showing at the end of this section the general kernel renormalization scheme for \(L\)-layer CNNs. In a 2HL CNN, the backward integration of the second layer proceeds in the same way as before, except with the replacement of \[\begin{align} & G \to \omega(K_E^{(1)})\, , \\ & [\omega(K_E^{(1)})]^{\mu\nu}_{i_1 i_1'} = \frac{1}{M} \sum_{m} [K_{E,S_2 i_1+m, S_2 i_1'+m}^{(1)}]^{\mu\nu} \, . \end{align}\] Here the function \(\omega(\cdot)\) is a simple linear function that maps the empirical kernel at layer \(\ell=1\) to the cross-covariance matrix of the pre-activations at layer \(\ell=2\). Note that, because of the different number of patches at each layer, while \(K_E^{(1)} \in \mathbb{R}^{N_{p_1}P \times N_{p_1} P}\), \(\omega(K_E^{(1)}) \in \mathbb{R}^{N_{p_2}P \times N_{p_2} P}\). We also note that, for the same reason, in 1HL CNNs we have (see Eq. 40 ): \[G = \omega(C)\, , \quad [C_{i_0 i_0'}]^{\mu\nu} = \frac{1}{A_0} \sum_{a_0}x^\mu_{i_0 a_0}x^\nu_{i_0' a_0} \, .\] In the equations above, we denoted with the same symbol \(\omega(\cdot)\) both the 1HL and the 2HL cases to keep the notation light, implying that, in principle, this function depends on the layer, since the number of patches varies across layers. From now on, the function \(\omega\) is used at every layer, implicitly encoding this layer dependence. This is a peculiar feature of CNNs, where the dimension of covariances between pre-activations and post-activations generally changes. Recalling Eq. 45 , in the 2HL case we consider the following integral: \[\begin{align} \varphi(\bar{f}|X) = \int_{\mathcal{S}_{N_{p_2}}^+} & dp(q_2) \int \prod_{a_1} dH^{(1)}_{a_1} \rho_\mathcal{N}\left(H^{(1)}_{a_1} ; 0, G \right) \times \nonumber \\ & e^{-\frac{1}{2N_{p_2}}\mathop{}\!\mathrm{Tr}[(\tilde{F}U_2)^\top \Theta(\omega(K_E^{(1)})) (\tilde{F} U_2)]}\, , \end{align}\] where \(q_2 \sim {\mathcal{W}_{N_{p_2}}}(\mathbb{1}/ A_2, A_2 )\) and in analogy with the multiple-output case, we used the decomposition \(q_2 = U_2 U_2^\top\). The stacked EWA for the matrix \(\Theta(\omega(K_E^{(1)})) = (\Theta\circ\omega)(K_E^{(1)})\) turns out to be: \[(\Theta\circ\omega)(K_E^{(1)}) \sim \mathcal{W}_{N_{p_2} P}\bigl( (\Theta \circ \omega )^2(C)/A_1, A_1 \bigr) \, ,\] which, using the same strategy as in Eqs. 42 , 43 , and 44 , gives \[\begin{align} \label{eq:2hl-cnn-final-integral} \varphi(\bar{f}|X) = \int_{\mathcal{S}_{N_{p_2}}^+} & \prod_{i=1,2} dq_i\;\rho_{\mathcal{W}_{N_{p_2}}}(q_i; \mathbb{1}/ A_i, A_i ) \nonumber \\ & \times e^{-\frac{1}{2N_{p_2}}\mathop{}\!\mathrm{Tr}[(\tilde{F}U_2)^\top (\Theta \circ \omega)^2(C) (\tilde{F} U_2) q_1]}\, . \end{align}\tag{46}\] Introducing the same decomposition, \(q_1 = U_1 U_1^\top\), the 2HL kernel renormalization reads: \[K^{(\text{R})}_{\mathcal{Q}} = \frac{1}{N_{p_2}}\sum_{ij=1}^{N_{p_2}}(\underbrace{U_2 U_1 U_1^\top U_2^\top}_{\mathcal{Q}})_{ij}(\Theta \circ \omega)^2(C)_{ij} \, .\] Given that the final 2HL integral in Eq. 46 has the same form as the 1HL integral, we can repeat the same procedure with the replacement \(\omega(K_E^{(1)}) \to \omega(K_E^{(2)})\) and apply the same machinery to compute the general \(L\)-layer case. This recursive scheme yields the following deep local kernel renormalization scheme: \[\begin{align} K^{(\text{R})}_{\mathcal{Q}} & = \frac{1}{N_{p_L}}\sum_{ij=1}^{N_{p_L}}\mathcal{Q}_{ij}\bigl[ (\Theta \circ \omega)^L(C)\bigl] _{ji} \, , \label{eq:cnn-kernel-renorm} \\ \mathcal{Q} & = \Biggl( \prod_{\ell=1}^{L}U_{\ell}^\top \Biggr)^\top\Biggl( \prod_{\ell=1}^{L}U_{\ell}^\top \Biggr) \, . \end{align}\tag{47}\]

6 Posterior sampling experiments at finite width compared to theory predictions↩︎

In this section we present an extensive comparison between the learning curves resulting from the EWA analysis and the outcomes of numerical simulations. The main purpose is to assess and quantify the predictive power of the EWA formalism in describing the behavior of fully-connected NNs as a function of the depth \(L\), the number of training patterns \(P\), and the number of neurons \(N_1,\ldots,N_L\). In contrast to Section 4.2, where the focus was on the rate function of \(\mathcal{Q}\) within the context of the characteristic of the prior, this numerical analysis centers on the posterior predictor statistics. It is worth noting that an additional numerical analysis is required at the level of the posterior, due to the impossibility of checking the EWA for each configuration of the dual outputs \(\bar{f}\) over which we integrate in Eq. 11 . In view of this, numerical checks at the posterior level provide a consistent probe of the capabilities of the EWA, taking into account all the approximations involved.

Here we focus on the training loss and the generalization error, as defined in Eq. 7 and Eq. 6 respectively. These two observables are expected to capture the fundamental behavior of the networks, probing both interpolation and generalization capabilities. Examples of posterior predictor distributions for individual test points can be found in supplementary Figs. 14-16. Since the numerical evaluation of Eqs. 7 and 6 requires sampling from high-dimensional and complex distributions, we employed various Markov Chain Monte Carlo (MCMC) techniques to approximate these ensemble averages via time-averaging over appropriate stochastic evolutions that can be implemented numerically. We provide the main details regarding the MCMC implementations in Sec. 6.2, while extended discussions are presented in App. 11.

6.1 Theory for predictor statistics↩︎

The theoretical predictions, on the other hand, are straight-forward to evaluate. Since the characteristic function of the prior in Eq. 25 is a Gaussian mixture in the variables \(\bar{f}\) and the mixture weight concentrates on \(q_1^*, \dots, q_L^*\) in the posterior, the predictor statistics given squared-error likelihood Eq. 3 are those of a Gaussian process regressor. Under these circumstances, analytical expressions for the train and test losses can be obtained by means of the bias–variance decomposition: \[\epsilon_{\text{g/t}}(x,y) = \left(y-\Gamma(x)\right)^2 + \sigma(x)\, ,\] where the bias and variance of the network outputs are those of a Gaussian process regressor [7] with the large-width NNGP kernel replaced by the renormalized kernel evaluated at the minimum of the effective action (see Eq. 34 ). Details regarding the numerical implementation employed to obtain the minimum \(q_1^*, \dots, q_L^*\) are provided in App. 13 (see also Ref. [63] for the GitHub repository). Thus, for a generic input \(x\) the bias and variance read: \[\begin{align} \tag{48} \Gamma(x) = & \left[ K^{(\text{R})}_{\mathcal{Q}^*}(Xx) \right]^T \left[ \frac{\mathbb{1}}{\beta} + K^{(\text{R})}_{\mathcal{Q}^*}(C)\right]^{-1} Y \, , \\ \tag{49} \sigma^2(x) = & \, K^{(\text{R})}_{\mathcal{Q}^*}(x^Tx)\;+ \left[ K^{(\text{R})}_{\mathcal{Q}^*}(Xx) \right]^T \nonumber \\ & \times \left[ \frac{\mathbb{1}}{\beta} + K^{(\text{R})}_{\mathcal{Q}^*}(C) \right]^{-1} \left[ K_{\mathcal{Q}^*}^{R}(Xx) \right] \, , \end{align}\] where \(Xx \in \mathbb{R}^{P}\), \(x^Tx \in \mathbb{R}\), \(\mathcal{Q}^*=\prod_{\ell=1}^{L}q_\ell^*\) and \[K^{(\text{R})}_{\mathcal{Q}^*}(Xx) = \mathcal{Q}^* \Theta^L(Xx)\, , \quad K^{(\text{R})}_{\mathcal{Q}^*}(x^Tx) = \mathcal{Q}^* \Theta^L(x^Tx)\, .\] In the equations above, \(\Theta^L(Xx) \in \mathbb{R}^P\) and \(\Theta^L(x^Tx) \in \mathbb{R}\) denote in analogy with Eq. 21 the \(L\)-th composition of the NNGP kernel function, acting respectively on a vector and on a scalar entry; specifically, \[\begin{align} [ \Theta(Xx) ]_\mu & = \mathbb{E}_{h^\mu, h}\left[ \sigma(h^\mu)\sigma(h)\right] / \lambda \nonumber \\ \Theta(x^Tx) & = \mathbb{E}_{h}\left[ \sigma(h)\sigma(h)\right]/\lambda \end{align}\] where \((h^\mu, h)\sim\mathcal{N}(0,\Sigma_\mu)\) and \(h\sim\mathcal{N}(0,\Sigma_s)\), being \(\Sigma_\mu = \frac{1}{N_0 \lambda} \left(\begin{smallmatrix} (x^\mu)^T x^\mu & (x^\mu)^T x\\ x^T x^\mu & x^Tx \end{smallmatrix} \right)\) and \(\Sigma_s = \frac{x^Tx}{N_0 \lambda}\). It is worth noting that in the limit \(\alpha \to 0\), the saddle-point solutions reduce to \(q_1^* = \ldots = q_L^* = 1\), and Eqs. 48 and 49 for the bias and variance reproduce the large-width asymptotics [8], as expected.

6.2 MCMC Sampling↩︎

In this work, we primarily employ the Langevin Monte Carlo (LMC) algorithm, a standard method for sampling from high-dimensional distributions that take the form of a Gibbs distribution. It is based on the continuous Langevin equation \[\label{eq:langevin-continuous} \dot{\theta}(t) = - \nabla_\theta \mathcal{L}_{\text{reg}}(\theta(t)) + \sqrt{2T} \epsilon(t) \, ,\tag{50}\] where \(T=1/\beta\) and \(\epsilon(t)\) is Gaussian white noise with moments \[\langle \epsilon(t)\rangle=0\, ,\quad \langle\epsilon(t)\epsilon(t')\rangle = \delta(t-t') \, .\] Indeed, the stationary probability distribution \(P_\beta(\theta) = \lim_{t \to \infty} P_{t}(\theta)\) defined by the dynamics in Eq. 50 , which results from averaging over all possible white noise realizations, is the posterior distribution \(P_\beta(\theta) \propto \exp{(-\beta \mathcal{L}_{\text{reg}})}\). In this work, we use the discretized version of Eq. 50 for the numerical implementation, namely: \[\label{eq:langevin-discrete} \theta_{t+1} = \theta_t -\eta\nabla_\theta\mathcal{L}_{\text{reg}}(\theta_t)+\sqrt{2T\eta}\epsilon_t\, ,\tag{51}\] where \(\langle\epsilon_t\rangle=0\), \(\langle\epsilon_t\epsilon_{t'}\rangle=\delta_{t,t'}\) and \(\eta\) is the step size for the numerical integration (or learning rate). Note that in Eqs. 50 and 51 , the gradients are computed using the regularized loss function: this function includes both the Squared Error likelihood contribution and the Gaussian term arising from the prior, i.e. \[\mathcal{L}_{\text{reg}}(\theta_t) = \mathcal{L}(\theta_t) - \frac{1}{\beta} \log{\rho(\theta_t)}\, .\] At different time scales, the stochastic evolution driven by Eq. 51 exhibits two distinct regimes. Close to the initialization, the dynamics is dominated by the gradient contribution over the white noise, as the system is initialized in a random configuration \(\theta_0\) that is typically far from equilibrium. This stage, known as the thermalization phase, persists for a small number of updates and is discarded from the statistical analysis presented below. At larger time scales, once the system reaches the typical configurations of the target distribution—for instance, after \(s\) updates—the equilibrium phase begins. This phase persists until the end of the simulation, and the configurations \(\{\theta_s, \theta_{s+1}, \ldots, \theta_{s+T}\}\) generated during this stage are used to approximate the ensemble average as: \[\langle B\rangle \approx \frac{1}{T} \sum_{t=s}^{s+T} B(\theta_t) \, ,\] where \(B(\theta_t)\) is a generic observable computed on the network configuration \(\theta_t\). Under general hypotheses, the Monte Carlo average converges to the ensamble average as the inverse square root of the number of samples, i.e., \(\langle B \rangle = \frac{1}{T} \sum_t B(\theta_t) + O(1/\sqrt{T})\). The rate of convergence depends on the characteristic time required by the algorithm to generate independent samples, known as the autocorrelation time [64]. To account for autocorrelation effects in the estimation of the statistical error, we computed the uncertainty of the mean using the blocking method (see App. 11.2 for further details). In this work, we generated chains of \(O(10^7)\) network samples for each experimental point, depending on the computational load at hand, which proved to be sufficient for a reliable estimation of both the mean and the autocorrelation time, and thus of the statistical uncertainty.

In addition to statistical errors, the discretized LMC also introduces systematic errors. Indeed, the stationary distribution of Eq. 51 is only an approximation of the Gibbs distribution, with an additional systematic bias arising from the numerical integration that scales linearly with the step size \(\eta\), i.e., \(P_\beta(\theta) \propto e^{-\beta\mathcal{L}(\theta)}\rho(\theta) + O(\eta)\). To mitigate finite step size effects, we used a small learning rate of the order \(\eta \sim 10^{-3}\) such that the systematic error remains smaller than the statistical uncertainty. We verified this condition a posteriori by further reducing the learning rate and showing that the new estimation of the mean is compatible with the previous one within our statistical precision (we report a representative example in Fig. 18).

It is worth noting that Eq. 51 corresponds to the standard Gradient Descent (GD) dynamics with the addition of a noise source proportional to the model temperature. Following this analogy, discretized LMC is a form of training dynamics that in the long term limit samples from the Bayesian posterior. Although sampling with MCMC algorithms is not the conventional approach for training modern deep networks due to their high computational cost, it represents a well-established numerical framework for investigating their behavior and performance (see for instance Refs. [7], [65]). Furthermore, outcomes from Bayesian inference at low temperatures \(T\) have been shown to be close to those obtained with GD [6], [44], [66], which further motivates the interest in such training algorithms.

6.3 Learning curves for fully-connected DNNs↩︎

We conducted extensive simulation campaigns to compare the predictions of the EWA with numerical outcomes obtained via Bayesian sampling. In particular, we measured both the training and generalization losses in DNNs using two stereotypic activation functions: (i) Erf, as a non-linear zero-mean activation function, and (ii) ReLU, as a non-linear activation function with mean. We focus on three different datasets: the MNIST dataset (classes “\(0\)” and “\(1\)”), the CIFAR-10 dataset (classes “cars” and “planes”), and a synthetic Gaussian dataset with linear teacher rule \(y^\mu=w^\ast x^\mu\). For MNIST and CIFAR-10, no one-hot encoding was applied to the labels, resulting in single-output DNNs. Additional details regarding the datasets are extensively discussed in App. 12.

In Figs. 1 and  4 we show the learning curves for DNNs with Erf activation functions, trained on the MNIST and CIFAR-10 dataset respectively. The first and the second rows report the training and test loss as function of the number of layers \(L\). The number of neurons is kept fixed at \(N_1 = \ldots = N_L = 200\), while the number of training patterns \(P\) varies across columns, such that different columns correspond to different values of \(\alpha_1 = \ldots = \alpha_L = \alpha\). In addition to the EWA predictions (green solid lines) and numerical simulations (black dots), each panel also reports the large-width NNGP prediction obtained at the same \(P\) and \(\alpha=0\) (gray dashed lines). In all panels, numerical simulations differ quantitatively from the large-width theory (and often also qualitatively, see especially Fig. 4) as the network depth increases. In contrast, the EWA predictions are in good agreement, even for a significant number of hidden layers. As \(\alpha\) increases, corresponding to the right columns in the figure, the large-width theory becomes increasingly inaccurate even for small depths, while the EWA continues to quantitatively track the experimental behavior, apart from small discrepancies due to the EWA approximation that nonetheless remain limited even for large depths. In the Appendix, Fig. 19, we show simulations for DNNs with Erf activations on the random Gaussian dataset. Again, the learning curves obtained with the EWA are predictive across the entire range of depths considered, especially for \(\alpha < 2.5\). In panel (h) (\(\alpha=2.5\)) we observe discrepancies between numerical outcomes and EWA predictions for larger depths; nevertheless, even in the worst cases, the EWA approximations for the generalization error exhibit a mismatch no larger than \(5\)–\(8\%\).

In Fig. 5 we show learning curves for DNNs with ReLU hidden units for the random Gaussian dataset. Here the parameter of the simulations are the same as Figs. 1 and 4, with the exception of the choice of the Gaussian prior precisions. While in the latter we used standard priors \(\lambda=1\), in the ReLU case we tune the precision to critical priors \(\lambda=1/2\) so that the relative NNGP kernels are well-behaved at large depth [67]. For ReLU the large-width theory again fails to qualitatively capture the behavior of the DNNs. As shown, in these regimes learning is completely dominated by finite-width effects which are neglected by the infinite-width theory. Instead, the EWA learning curves do successfully describe the behavior of ReLU DNNs at finite width. In addition, in Fig. 5 (f, g), the large-width and EWA learning curves cross around \(L=3\), with numerical simulations consistently matching the EWA. This indicates that our working regime is far from one where finite-width effects can be treated as mere first-order perturbative corrections. In such regimes, first-order corrections to an observable \(A\) are expressed as \(A = A^{(0)} + \alpha L A^{(1)} + \dots\), where \(A^{(0)}\) is the Gaussian infinite-width prediction and \(A^{(1)}\) is a constant first-order perturbative correction [67]. Since learning curves that take into account only first-order corrections to the Gaussian large-width theory can correspondingly never cross the zero-order prediction as a function of depth, our numerics clearly shows that the EWA accounts for strong, non-perturbative finite-width effects in a layer-dependent manner. We conclude this section noting that in Fig. 5 (h) in analogy to Fig. 19 (h), the EWA starts to break down for ReLU at \(\alpha=2.5\) and depths \(L>3\), when the network is trained on random Gaussian inputs. These deviations are expected to gradually emerge as the network depth increases, reflecting the approximate nature of the EWA. Similar effects have already been discussed in recent literature, where in the Bayes-optimal setting of student MLPs trained to match random teacher networks, an upper-bound for the generalization can be rigorously proved, proceeding by a layer-wise reduction [22]. The upper-bound scales as \(O(L/\sqrt{N})\), accumulating errors for each layer reduction. While in our setting with general data a different scaling with \(N\) and \(P\) is possible, we expect the same scaling of deviations with \(L\) holds due to the similar layer-wise structure of the computation. Interestingly, for simulations with MNIST and CIFAR-10 the EWA remains predictive also at \(\alpha=2.5\) and \(L>3\) (see panels (d, h) of Appendix Figs. 20 and  21). The same scaling argument nevertheless suggests that consistent deviations could still appear at larger depths than those explored here.

6.4 Emergent metastable regime at \(L\alpha\gg 1\)↩︎

While the EWA learning curves provide an overall proper description of the behavior of NNs across a broad range of network architectures and learning tasks, in the regime where both \(L\) and \(\alpha\) are sufficiently and simultaneously large, numerical simulations reveal sudden improvements in performance. In particular, we find a sharp reduction in both the train loss and the generalization error in some cases as depth increases, unveiling an emergent and potentially new form of finite-width contributions. In this section, we provide a systematic investigation of this phenomenon, relating these empirical observations to a metastability in the Bayesian learning dynamics, and perform a finite-size study against the EWA predictions.

In Fig. 6 we show MCMC experiments on MLPs with ReLU nonlinear hidden units trained on the same CIFAR-10 task as in Figs. 4 and 21. All the simulations in this case refer to the temperature \(T=0.01\), \(\alpha = 2.5\), and critical priors. Panels (a) and (b) report the main message: at \(L=6\), both the train and the test loss exhibit sudden drops (see colored circles). It is worth mentioning that these numerical outcomes are obtained using the No-U-Turn Sampler Hamiltonian Monte Carlo (NUTS HMC) [68], a state-of-the-art MCMC sampling algorithm that ensures improved performance compared to LMC. While being built on top of the vanilla HMC, NUTS HMC integrates features such as automatic fine-tuning of hyperparameters, including the integration step and trajectory length. In our case, the more powerful but less interpretable NUTS HMC is implemented to avoid finite learning-rate systematics and counteract slowdowns in the Bayesian sampling at large \(L\alpha\). Posterior sampling experiments clearly show that: (i) the transition is not a finite-size effect that vanishes in the proportional regime \(N_\ell, P \to \infty\), since the drop becomes more and more pronounced as we simulate networks with a larger number of neurons and training patterns (see the different colored circles in panels (a, b), which show increasing \(N_\ell\) and \(P\) at fixed \(\alpha\)); (ii) both the large-width theory and the EWA continue to predict smooth behavior in these regimes. Since consistent Bayesian sampling is known to be challenging under such conditions, i.e. where one has simultaneously large depth, a large number of training examples, and low temperature, we implemented several MCMC samplers to ensure the robustness of our results. In panel (c) of Fig. 6 (shaded points) we also display results obtained with the Metropolis-Adjusted Langevin algorithm (MALA) [69], roberts1996exponential? and the preconditioned Crank–Nicolson MCMC algorithm (pCN) cotter2013mcmc?, along with LMC and NUTS HMC results. We highlight that while LMC and MALA are gradient-based samplers, pCN is an energy-based sampler, whereas NUTS HMC is a momentum-driven sampler, so that biases arising from common strategies in exploring the configuration space are avoided. The results of the numerical experiments are fully in agreement with each other, providing strong numerical evidence that additional finite-width effects in DNNs are at play.

Quite interestingly, we found that the Bayesian sampling dynamics exhibit for those points at \(L>5\) a new intermediate phase between the thermalization and the equilibrium phases, which we refer to as a transient phase. Configuration updates in the transient phase are characterized by a balance between gradient forces and random fluctuations, so that the system seems to be at thermal equilibrium for long simulation times, until the thermalization dynamics restart and drive the system toward the true long-term equilibrium distribution. The number of epochs needed to escape the transient phase depends on the algorithm and can last up to millions of updates, and therefore may persist longer than the sampling experiment even for simulations with a large number of epochs. An example comparing two independently initialized LMC chains is shown in Fig. 15. It is worth mentioning that in our experience we observed drifts from the transient to the equilibrium phase, and never reverse transitions. Taking advantage of these equilibrium-like properties of the transient phase, we computed the ensemble averages over those configurations and found that they not only smoothly continue the empirical learning trends observed at smaller depths \(L\le5\), but are also in agreement with the smooth behavior predicted by the EWA theory. In Fig. 6, panels (d, e), we display the onset of the transient phase in NUTS HMC for the train and test losses. Even if for NUTS HMC the transient phase lasts a few hundred epochs, due to the optimal preconditioning of the simulation parameters and much larger step distance per sample, it is evident that the EWA predictions (green lines) are in agreement with the Bayesian sampling, while after the drop the system fluctuates around different loss minima (the red line indicates the average over the equilibrium phase). In Fig. 6, panel (c), we show all the different averages together: shaded points correspond to equilibrium averages, while unshaded points represent transient averages. Up to \(L=5\) no drops are detected, while for \(L \ge 6\) we report both equilibrium and transient measurements, with the latter being in good agreement with the smooth learning curves predicted by the EWA.

6.5 Learning curves for shallow convolutional nets↩︎

In Sec. 5.2 we introduced the stacked EWA in order to describe the finite-width behavior of deep CNNs. For these architectures, the EWA requires the introduction of the stacked NNGP kernel \([(\Theta\circ\omega)^L(C)_{ij}]^{\mu\nu}\), which quantifies, in the large-width limit, the effect of nonlinearities when applied to patch–patch correlated pre-activations. Nevertheless, patch–patch correlations are discarded in the infinite-width limit, because only the diagonal components of the stacked NNGP kernel enter the predictor statistics [62]. In the proportional regime, the stacked EWA takes into account these correlations, which are captured by the low-dimensional emerging order parameter \(\mathcal{Q} \in \mathbb{R}^{N_{P_L} \times N_{P_L}}\) (note that the dimensionality of the order parameter depends only on the number of patches at the last layer). It is important to highlight that, after the dimensional reduction, the characteristic function of the prior can again be reduced to a Gaussian mixture in the quadratic variables \(\bar{f}\), which leads to the same expression for the predictor statistics in Eqs. 48 and 49 , only replacing the saddle-point value of the renormalized kernel for MLPs with the CNN one, Eq. 47 . Here, the saddle-point approximation of the effective action is computed over the ensemble of \(N_{P_L} \times N_{P_L}\) positive-definite matrices rather than on scalar order parameters, allowing patch–patch correlations to contribute to the bias and variance. In Fig. 7 we quantify the predictive power of the stacked EWA by comparing it against Bayesian sampling experiments. In particular, we consider two settings of mask, stride, and model temperature \(T\) (different columns). Using Eqs. 48 and 49 , we compute the train and test loss (different rows) of a 1HL Erf CNN trained on MNIST, as a function of the number of channels \(A_1\) in the first layer at fixed \(P = 500\). As expected, the EWA predictions (green solid lines) approach the large-width ones for \(A_1 \gg P\) (dashed gray lines). At large but finite \(A_1 \sim P\), numerical simulations with real CNNs display a finite-width behavior that is perfectly tracked by the EWA learning curves, thus establishing the validity of the ansatz in describing architectures beyond fully connected DNNs. For additional details regarding the computation of the saddle point of the effective action for CNNs, we refer the reader to App. 13.

6.6 Mean field versus Standard parametrization of the weights↩︎

Here, we ask in how far the low-dimensional effective action obtained through the EWA can also be valid in the rich learning regime. To do so, we change from standard parametrization to the \(\mu\)P by defining \(f_{\mu\mathrm{P}} = \frac{1}{\gamma} f_{\mathrm{SP}}\), where \(\gamma = \gamma_0 \sqrt{N}\). From the Bayesian perspective, this corresponds to reducing the scale of the prior by factor \(\gamma\) and is equivalent to reducing the the variance of the weights in the readout layer as \(\lambda_L^{-1} \to \gamma^{-2} \lambda_L^{-1}\). Simply changing \(\lambda_L\) is convenient in the theoretical expressions. However, in the gradient-based LMC simulations we use instead the explicit factor in the network output, \(\frac{1}{\gamma} f_{\mathrm{SP}}\), which corresponds to the same output posterior \(P_\beta(f)\) but is a choice of weight parameters in which gradients in all layers are reduced equally, allowing to use the standard increase of learning rate in \(\mu\)P of \(\eta \to \gamma^2 \eta\).

As a final consequence of the change from SP to \(\mu\)P, also the temperature has to be scaled down as \(\frac{1}{\gamma^2} T\) to avoid dominance of the sharper prior over the likelihood term [52], [53]. This means that SP and \(\mu\)P can not be compared consistently at finite temperatures, since \(\mu\)P requires \(T\to 0\) to perform better than samples from the prior in the high-dimensional limit. This susceptibility to noise can be seen as a limitation of the \(\mu\)P parametrization. However, the zero-temperature limit is not unrealistic since in practical training, output noise due to floating point error is small and noise arising from finite batch sizes and learning rates is progressively reduced over training time. At the finite sizes of \(N\sim 10^3\) that we probe in our experiment, we found that using a temperature \(T=10^{-1}\gamma^2 \sim 10^{-4}\) was both feasible to sample and sufficient to obtain behavior qualitatively similar to that in the zero-temperature limit.

Overall, we find that the EWA successfully predicted generalization performance even in the rich learning regime. Fig. 8 shows a comparison of SP (green) and \(\mu\)P (blue) for ReLU MLPs of depth \(L=4\) trained on the CIFAR-10 “cars” vs. “planes” task with \(P=1000\). The behavior of the generalization error is well captured in both parametrizations, and again clearly differs from the large-width limit (dashed gray). In panels (b,c), the error is decomposed into the contributions arising from the bias and variance of the posterior predictor distribution. The substantial improvement of the performance in \(\mu\)P seen in (a) is almost entirely explained by a reduction of variance in \(\mu\)P (c), while the bias remains similar to SP (b).

It appears counter-intuitive that the EWA can capture the behavior of \(\mu\)P, because the renormalized kernel Eq. 33 does not adapt in the same way as the empirical hidden-layer kernels do in the rich regime. It is a characteristic of \(\mu\)P that on a binary classification task as shown in Fig. 8, the posterior empirical kernels show a strong block-structure when the sample indices are sorted by class label; while the renormalized kernel retains the prior structure and only receives a global rescaling factor. However, this rescaling could effectively take into account the adaptation of the output posterior to the data. Indeed, Fig. 8 (e) shows that the empirical kernel of the last-layer representation has acquired a block structure consistent with strong feature learning of the \(\{0,1\}\) class structure, while the generalization error remains well predicted by the unstructured effective kernel (d). This at first surprising behavior is consistent with deep linear networks however. Here it can be shown that the renormalized effective kernel is equivalent to adding a learned low-rank contribution \(\propto ff^\top\) to the empirical kernels at each layer [23], [24], which corresponds to a block structure like that observed in the empirical kernel undergoing feature learning.

For the single output task as in Fig. 8, in the zero temperature limit the renormalization factor of the kernel cancels out in the mean predictor Eq. 48 while entering linearly into the predictor variance Eq. 49 , which is strongly reduced in \(\mu\)P. The effect is that at low temperature, the improvement of MLPs in the rich learning regime over SP is mostly attributable to a difference in variance, while the biases are similar when the EWA is a valid approximation. This can explain the behavior observed in the experiments on real data we performed in the proportional sample-width regime. However, this does not preclude the existence of tasks where the EWA breaks down more strongly and a qualitative difference in the biases appears. A detailed empirical and theoretical exploration of output posteriors \(P_\beta(f)\) in the rich learning regime will therefore be the subject of a follow-up work.

7 Discussion and Conclusions↩︎

We have presented an approximate theory describing the generalization error in Bayesian multi-layer, nonlinear networks at finite width. In the proportional scaling regime where the number of samples \(P\) and the width \(N\) are comparably large, this theory successfully captures the empirical performance of Monte-Carlo sampling experiments across multiple architectures and datasets, particularly for moderate network load \(\alpha = \frac{P}{N} \lesssim 1\).

Building upon work on deep linear networks, we show how an Equivalent Wishart Ansatz reduces the intractable hierarchy of empirical kernel fluctuations to a low-dimensional effective theory in terms of a renormalized kernel and a small number of self-consistent order parameters. Within this framework, both fully-connected and convolutional layers can be treated, and we expect that generalizations to other architectures admitting a well defined large-width limit are possible.

A few additional aspects are worth emphasizing: First, the effect of non-zero activation means is asymptotically suppressed, so that the zero-mean EWA remains the relevant effective description also for activation functions such as ReLU. Second, the theoretical predictions are supported by an extensive set of Bayesian sampling experiments across depths, datasets, and activation functions, which to our knowledge constitutes the broadest such test so far performed on deep MLPs. Finally, the results suggest that learning regimes in large but finite-width networks can be organized by the network load \(\alpha=P/N\), and therefore by the relative scaling of data and width, rather than only by the usual lazy-versus-rich parametrization distinction.

7.1 Focusing on the layer-wise kernels as random matrix ensembles↩︎

In this paper, we focus on the kernel matrices as the fundamental random variables of the learning system. A theory describing deep network learning outcomes in a Bayesian setting is then always a more or less faithful approximation of the prior ensemble of kernel matrices. In the EWA, we aggregate all randomness in the lower-layer kernels into approximately Wishart fluctuations of the last-layer kernel matrix. This notably differs from a Gaussian approximation in the pre-activations in each layer, which would give the infinite-width NNGP result.

Since a Wishart kernel by definition in Eq. 8 also follows from Gaussian fluctuations in the post-activations, a Gaussian Ansatz for the last-layer post-activation distributions arising from each of the lower layer weight priors would also imply the EWA. However, this transports no physical insight, since the individual nonlinear post-activations in reality are never Gaussian distributed, also not in the infinite-width limit. The kernel matrix entries (suppressing layer indices for brevity) \(K^{\mu \nu}=\frac{1}{\lambda N} \sum_i \sigma(h^\mu_i)\sigma(h^\nu_i)\) instead, being always high-dimensional sums of products, behave close to Wishart even if the \(\sigma(h_i)\) are not Gaussian, and also the continuation to the infinite-width limit holds. The perspective of approximating the kernel distribution is therefore physically more natural. It is nonetheless important to note that it is still an equivalent Wishart Ansatz: For the EWA to hold, the kernel needs not be exactly Wishart but we only need the implication \(Q=\frac{\bar f^\top K \bar f}{\bar f^\top \mathbb{E}[K] \bar f} \sim \chi^2_N\) to hold for those \(\bar f\) which dominate the partition function Eq. 11 .

Overall, we argue that the ensemble of kernel matrices may be the more convenient object to study learning in these systems. In deep linear networks, a full understanding of the posterior appeared as a highly complex problem [26], [39] until by considering the ensemble of kernel matrices the posterior distribution could be characterized exactly at arbitrary values of \(N,P\) and \(L\) [24]. With this framing as a random matrix problem we hope to stimulate also mathematically rigorous work on the kernel ensembles in nonlinear multi-layer networks.

It may appear as a limitation that pairwise kernels are naturally sufficient statistics for square loss, while cross-entropy loss takes also higher-order information of the output prior into account. However, we note that the joint output prior Eq. 16 can always be expressed as an integral representation in terms of the pairwise kernel. Using cross-entropy loss only means that the posterior is no longer an analytic Gaussian integral of the prior characteristic function.

7.2 Finite- versus infinite-width generalization performance↩︎

The proportional sample/width limit provides a natural framework to study large but finite-width overparametrized networks. Throughout, we have seen that the finite-width generalization performance can differ strongly from the infinite-width (lazy) NNGP prediction.

At the same time, our experiments and approximate theory show that for MLPs the finite-width correction takes a surprisingly simple form: to leading order, the effect is captured by a scalar renormalization of the kernel Eq. 33 by a factor \(\mathcal{Q}\). This indicates that at network load \(\alpha = O(1)\) the amount of data structure extracted by deep fully-connected networks is limited. The significant difference to the infinite-width limit \(\alpha \to 0\) is here driven mostly by a difference in the predictor variance. This is in contrast to deep convolutional architectures, where the local kernel renormalization Eq. 47 is naturally more structured. Because of this, at \(\alpha = O(1)\) also the bias of finite-width CNNs can differ substantially from the infinite-width limit. From a broader perspective, while much theoretical work has focused on the rich learning regime to go beyond the lazy infinite width limit, we argue that the relative scaling of the number of samples and the width, encoded by \(\alpha\), may be a more fundamental organizing principle for finite-width learning than the lazy-versus-rich dichotomy formulated only in terms of output scaling.

7.3 EWA as a baseline: Relation to adaptive kernel theories↩︎

A line of recent works have studied feature learning in the mean-field or \(\mu\)P setting, by deriving theories in which the full \(P\times P\) kernels are high-dimensional order parameters adapting to the data in the posterior[50]–[52] or during gradient descent dynamics [52], [70], typically together with backward-propagating conjugate or gradient kernels that are also fully adaptive. These are valuable and complex frameworks, but they are still intrinsically high-dimensional: the dominant saddle point must be found in coupled \(P\times P\) matrix equations whose complexity, in the proportional regime, is not fundamentally simpler than that of sampling the original \(N\times N\) weight posteriors. Ref. [28] provides a unifying account at the example of shallow networks, where the dimensionality of the order parameters can furthermore be reduced to \(P\)-dimensional mean-output discrepancies. In Ref. [71], steps are being made to simplify the \(L\)-layer \(P\times P\) equations by choosing a lower-dimensional variational Ansatz, which we believe is a very promising future direction. Another notable approach is Ref. [53], which also requires high dimensional optimization but focuses on nonlinearity-dependent inhomogeneities in the activation prior, that can induce discrete or sparse coding schemes in the activation posterior if outputs are further scaled down compared to \(\mu\)P.

Against this background, the EWA provides a strong and missing baseline. It is a direct nonlinear generalization of the type of kernel adaptation that is expected in deep linear networks, and remains easy to compute and interpret for deep architectures. In the present manuscript this lets us study networks with ten hidden layers, going beyond the two- to four hidden layer networks presented in previous adaptive-kernel studies. Moreover, the results show that especially in the regimes at moderate depth and load which are most accessible to sampling experiments, a much simpler theory already explains learning and generalization with high accuracy. For this reason, we believe that the EWA should serve as a baseline when studying the emergence of more directional forms of feature learning in nonlinear networks: comparing only to the infinite-width limit can be misleading, because the large gap between infinite-width NNGP and finite-width phenomenology is mostly captured by the directionally homogenous renormalized-kernel picture.

7.4 Emergence of deviations from the EWA theory at large depth \(L\) and large load \(\alpha\)↩︎

Aside from the good fit obtained in the majority of the experiments, we observed two types of systematic deviations from the EWA, both appearing when the depth \(L\) and the load \(\alpha\) become large. In both cases we performed additional sampling experiments, which suggest that the effect is not due to a simple failure or bias of the sampler, even though sampling itself becomes more difficult and has longer convergence times in this regime.

The first type of deviation develops gradually with increasing depth in the Gaussian-data experiments at \(\alpha=2.5\), see Figs. 5 (h) and 19 (h). There, the EWA approximation appears to deteriorate smoothly as \(L\) grows, while notably the agreement remains very good at \(L=1\). This could reflect an accumulating error in the predicted width of the rate function of the product \(\prod_\ell q_\ell\), or it could be a sign of directional inhomogeneity in the prior that is neglected by the EWA. Furthermore, also finite-size effects are expected to accumulate across layers in deep networks [67], [72], [73].

The second type of deviation is qualitatively different. For the CIFAR-10 task with ReLU activations, around \(\alpha=2.5\) and \(L>5\), we observed an abrupt transition in both training and test error (Fig. 6 and Fig. 15). This phenomenon is interesting and distinct from the gradual drift described above, showing all the signatures of a metastability in the posterior where the macro-state described by the EWA loses stability. It is not a finite learning-rate effect of the Langevin sampler: we replicated it with several sampling algorithms that reduce or eliminate finite-step-size bias, and found that increasing \(N\) and \(P\) sharpens the transition. At present we are not aware of a known mechanism in the literature that could explain this effect. Both types of deviations therefore point to not yet understood phenomena that arise specifically in deep and nonlinear networks, rather than shallow ones, and they provide a natural motivation for further theory development. We have begun follow-up investigations in this direction.

7.5 Sampling from posteriors with millions of dimensions↩︎

Due to the curse of dimensionality, guarantees on equilibration times for MCMC samplers become uninformative already at \(d\gtrsim100\) and particularly in the very high dimensions relevant for deep learning. Before the seminal work of Li and Sompolinsky on deep linear networks [23], it was therefore often viewed as beyond the capabilities of MCMC methods to reliably sample from the posterior over the weights of deep neural networks. This skepticism was also tied to the older picture of deep network loss landscapes as being filled with bad local minima that would only be avoided by dynamical biases of the optimization algorithms, a view that has since been revised substantially in overparametrized settings: these loss landscapes have relatively flat, non-convex basins connected to the global minima, instead of a proliferation of local minima trapping dynamics [74]–[77].

Here we provided a systematic study showing to which extent MCMC sampling is reliable also for deep and nonlinear networks, across several datasets, activation functions, and prior parametrizations. In the experiments presented here, even at system sizes with \(L \times N^2 \sim 10^6\), independent Monte-Carlo chains produce consistent posterior statistics over feasible simulation times, including predictor distributions resolved at the level of individual test points (Fig. 14). This can partly be explained by the relevant function-space observables being only the \(O(P)\) output degrees of freedom. Overall, scaling up the network size and the load parameter \(\alpha\) slows down sampling, yet not so severely that at \(\alpha \approx 1\) system sizes of \(N,P \sim 10^4\) or more would be out of reach.

Nonetheless, exponentially long timescales can be hidden behind apparently flat histories, such that sampling experiments should be interpreted with restraint. Indeed, the metastability phenomenon we report in Figs. 6, 15 happened for \(L=6\) on a timescale of millions of steps for LMC, MALA, and pCN with all predictor statistics seemingly converged. The available evidence thus suggests that large-scale Bayesian sampling in deep networks is feasible enough to be scientifically useful, but also subtle enough that equilibration can not be taken for granted.

7.6 Beyond the proportional regime↩︎

Our results point to the relative scaling of sample size \(P\) and width \(N\) as the key control parameter for the complexity of the output prior, and therefore the posterior adaptation structure. In the proportional regime, already the large-deviations structure of the prior becomes important. For MLPs the adaptation captured by the EWA remains low-dimensional, while successfully explaining the empirical learning curves seen in our sampling results.

This suggests that genuinely high-dimensional adaptation in MLPs may require scaling regimes beyond \(P \sim N\), where the number of samples grows fast enough to probe richer posterior structure than the one encoded by a renormalized kernel alone. This perspective is consistent with recent work on random teacher-student problems for deep MLPs, where a quadratic number of samples is needed for specialization on the teacher function [21], [78]–[80]. The proportional, overparametrized regime is a good proxy for some practical machine learning tasks, while other cases leveraging extreme amounts of data, especially for pre-training of vision- or language models [81] could be better described by quadratic scaling, falling outside the overparametrized regime. Understanding the capabilities and limitations of the most common deep learning architectures as a function of the sample size \(P\) relative to the network’s width \(N\) is therefore an important direction for future work.

7.7 Outlook↩︎

In this work we have proposed a tractable effective theory for Bayesian deep networks in the proportional regime, showing that a renormalized kernel description can capture finite-width learning curves in nonlinear MLPs far beyond the lazy infinite-width limit and on real-world datasets.

Several directions now appear natural. A first step is to extend the framework to more elaborate convolutional architectures including max-pooling, and to other network families with a richer internal geometry. Secondly, further scaling up sampling experiments could both enable direct comparisons with benchmark performance values and to map systematically if and where the EWA baseline deteriorates as depth and network load increase. Understanding those deviations, and determining whether they reflect finite-size effects or missing ingredients in the effective description, should be a particularly interesting direction for future work towards a theory of learning in deep networks.

8 Acknowledgements↩︎

R.B., R.P., P.R. and P.B. are supported by \(\#\)NEXTGENERATIONEU (NGEU). R.B. and P.R. are funded by the Ministry of University and Research (MUR), National Recovery and Resilience Plan (NRRP), project MNESYS (PE0000006) “A Multiscale integrated approach to the study of the nervous system in health and disease” (DN. 1553 11.10.2022). R.P. and P.B. are funded by MUR project PRIN 2022HSKLK9 and P2022A889F. This research benefits from the HPC facility of the University of Parma.

9 Non-central equivalent Wishart Ansatz↩︎

9.1 Contractions of the non-central Wishart distribution and definition of \(\tilde{Q}_\ell\)↩︎

The non-central Wishart distribution is obtained by allowing a non-zero mean in the Gaussian vectors entering the Wishart construction. Let \(G\) be a \(P \times N\) matrix, whose columns \(G_i\) (\(i=1,\dots,N\)) are independent and distributed as \[G_i \sim \mathcal{N}_P(\mu_i,V)\,,\] with common covariance matrix \(V\) and possibly non-zero means \(\mu_i \in \mathbb{R}^P\). Denoting by \(M = (\mu_1,\dots,\mu_N) \in \mathbb{R}^{P\times N}\) the corresponding mean matrix, the random matrix \(G G^\top\) is said to follow a non-central Wishart distribution, which we denote by \[R = G G^\top \sim \mathcal{W}^\mathrm{nc}_P(M M^\top,V,N)\,,\] with non-centrality encoded by \(M M^\top\). For \(M=0\) one recovers the ordinary central Wishart distribution. Another convention frequently used in the literature is to define the non-centrality matrix as \(\Omega = V^{-1} M M^\top\).

The only property of the non-central Wishart distribution that we will need is the analogue of Eq. 10 . Given a fixed vector \(s \in \mathbb{R}^P\), one has \[\frac{s^\top R s}{s^\top V s} = \sum_{i=1}^N (z_i)^2, \quad \quad \mathrm{where}\; z_i =\frac{s^\top G_i}{\sqrt{s^\top V s}}\, \sim \mathcal{N}\Big(\frac{s^\top \mu_i}{\sqrt{s^\top V s}}, 1\Big).\] Since the projected variables \(z_i\) are independent Gaussian with unit variance, it follows that the normalized contraction is distributed as a non-central chi-squared random variable: \[\frac{s^\top R s}{s^\top V s} \sim \chi_\mathrm{nc}^{2}(N, N\lambda_s)\,, \quad \quad \mathrm{where}\; N \lambda_s = \sum_{i=1}^N \frac{(s^\top \mu_i)^2}{s^\top V s} = \frac{s^\top M M^\top s}{s^\top V s}\,. \label{eq:property95chi95noncentral}\tag{52}\] In the special case relevant to our setting, in which all columns have the same mean \(\mu_i=m\), the non-centrality parameter reduces to \[N \lambda_s = N\,\frac{(s^\top m)^2}{s^\top V s}\,.\]

Definition of \(Q_\ell\) variables in the non-central EWA↩︎

Analogously to Eq. 24 we can introduce the variables \(Q_\ell\) \[\begin{align} e^{-\frac{1}{2} \bar{f}^\top K_E^{(L)} \bar{f}} &= \exp\bigg[ -\frac{1}{2} \underbrace{ \frac{\bar{f}^\top K_E^{(L)} \bar{f}}{\bar{f}^\top \Theta(K_E^{(L-1)}) \bar{f}}\; }_{\coloneq Q_L/N_L} \bar f^\top \Theta(K_E^{(L-1)}) \bar f \bigg] \nonumber \\ =\dots&= \exp\left[ -\frac{1}{2} \bigg(\prod_{\ell=1}^L \frac{Q_\ell}{N_\ell} \bigg) \bar{f}^\top \Theta^L(C_X) \bar{f} \right]. \end{align}\] For brevity we now use the short hands \(\Sigma_{(\ell)},m_{(\ell)}\) and \(\Theta_{(\ell)}\) instead of \(\Sigma\big(\Theta^{L-\ell-1}(K^{\ell-1}_E)\big)\), \(m\big(\Theta^{L-\ell-1}(K^{\ell-1}_E)\big)\) and \(\Theta^{L-\ell}\big(K^{(\ell-1)}_E\big)\). Given the non-central EWA \[\Theta^{L-\ell}\big(K_E^{(\ell)}\big) | K_E^{(\ell-1)} \sim \mathcal{W}^\mathrm{nc}_P \left(N_\ell\, m_{(\ell)} , \Sigma_{(\ell)}, N_\ell \right)\] the property Eq. 52 holds for contractions normalized by the scale matrix, which for non-zero means \(m_{(\ell)}\) is the covariance kernel \(\Sigma_{(\ell)}\) and not the second moment \(\Theta_{(\ell)} = \Sigma_{(\ell)} + m_{(\ell)} m_{(\ell)}^\top\) appearing in the definition of \(Q_\ell\). This is why we introduce the related variables \(\tilde{Q}\) in the non-central case \[\begin{align} \label{eq:tildeQ} \tilde{Q}_\ell &\coloneq \frac{\bar{f}^\top \Theta^{L-\ell}\big(K_E^{(\ell)}\big) \bar{f}}{\frac{1}{N_\ell}\bar{f}^\top \Sigma_{(\ell)} \bar{f}} \nonumber \\ &= \frac{\bar{f}^\top \big(\Sigma_{(\ell)} + m_{(\ell)} m_{(\ell)}^\top \big)\bar{f}}{\bar{f}^\top \Sigma_{(\ell)} \bar{f}} \; \frac{\bar{f}^\top \Theta^{L-\ell}\big(K_E^{(\ell)}\big) \bar{f}}{\frac{1}{N_\ell}\bar{f}^\top \Theta_{(\ell)} \bar{f}} \nonumber \\ &= \big(1 + \lambda_{\mathrm{nc},\ell} \big) Q_{\ell} \end{align}\tag{53}\] which according to Eq. 52 follow \(\tilde{Q}_{\ell} \sim \chi_\mathrm{nc}^2\big(N_{\ell}, N_{\ell}\,\lambda_{\mathrm{nc},\ell}\big)\) with non-centrality \(N_{\ell}\,\lambda_{\mathrm{nc},\ell} = N_{\ell}\,\frac{\bar{f}^\top m_{(\ell)} m_{(\ell)}^\top \bar{f}}{\bar{f}^\top \Sigma_{(\ell)} \bar{f}}\).

With these definitions one can proceed with the arguments given in Section 4.1 to show the asymptotic equivalence of central and non-central EWA. Note however, that if one wanted to write out an explicit effective action in the \(L\)-layer non-central case, this would require to track dependencies between the \(\tilde{Q}_\ell\) (or equivalently \(Q_\ell\)) distributions, which through \(\lambda_{\mathrm{nc},\ell}\) now explicitly depend on \(\bar f\) and \(m_{(\ell)}\) and therefore require additional order parameters to decouple, unlike in the central case where the \(Q_\ell\) distributions are independent and identical.

9.2 Non-central EWA for one hidden layer↩︎

For one hidden layer, the action arising from a non-central EWA can be written out without too much complication. This appendix analyses the relation to the central EWA result, showing how task-kernel and task-kernel-mean overlaps control the saddle point solution, and how the distributions of \(Q\) and \(Q_\mathrm{nc}\) differ while such a difference is absent in the output distribution.

We derive the action by first computing the characteristic function of the output prior Eq. 16 for a one hidden layer network under the non-central EWA, with the aim to plug into Eq. 11 and obtain the posterior partition function by a Gaussian integration. Recall the definitions of the kernels \(C = \frac{1}{N_0 \lambda_0} X X^\top\) and \(K^{\mu \nu}_\mathrm{E} = \frac{1}{N_1 \lambda_1}\sum_i^{N_1} \sigma_i^\mu \sigma_i^\nu\) and \(\Theta(C) = (\Sigma + m m^\top) / \lambda_1\); with hidden layer activations \(\sigma_i^\mu\) i.i.d. across \(i\), of mean \(m^\mu = \mathbb{E}[\sigma^\mu_i]\) and covariance \(\Sigma^{\mu \nu} = \mathrm{Cov}[\sigma^\mu_i \sigma^\nu_i]\). The characteristic function of the output prior Eq. 16 is \[\begin{align} \varphi (\bar f| X) & = \int_{\mathcal{S}^+_P} dK_{\mathrm{E}} \rho (K_{\mathrm{E}}|C) e^{-\frac{1}{2}\bar f^\top K_E\bar f} = \mathbb{E}\left[ e^{-\frac{1}{2}\bar f^\top K_E\bar f} \right] \\ & = \mathbb{E}\bigg[ \exp{\bigg( -\frac{1}{2} \underbrace{\frac{\bar f^\top K_\mathrm{E}\bar f}{\frac{1}{N_1 \lambda_1} \bar f^\top \Sigma \bar f}} _{\tilde{Q} \,\sim \,\chi^2_\mathrm{nc}(N_1, N_1 \lambda_\mathrm{nc})} \;\frac{1}{N_1 \lambda_1} \bar f^\top \Sigma \bar f \bigg)}\bigg] \end{align}\] Here \(\tilde{Q} \,\sim \,\chi^2_\mathrm{nc}(N_1, N_1 \lambda_\mathrm{nc})\) with non-centrality parameter \(\lambda_\mathrm{nc} = \frac{\bar f^\top m m^\top \bar f}{\bar f^\top \Sigma \bar f}\) under the non-central EWA. This can be seen by decomposing \(\tilde{Q} = \sum_i a_i a_i\) with \(a_i = (\bar f^\top \sigma_i) \big/ \sqrt{\bar f^\top \Sigma \bar f}\) and noting that \(\mathbb{E}[a_i] = \bar f^\top m \big/ \sqrt{\bar f^\top \Sigma \bar f}\).

Now expressing the distribution \(\rho(\tilde{Q}) = \int \frac{d\tilde{\bar Q}}{2\pi} \varphi(\tilde{\bar Q}) e^{-i\tilde{\bar Q} \tilde{Q}}\) via its characteristic function \[\varphi (\tilde{Q}) = \big(1 - 2i \tilde{\bar Q}\big)^{-\frac{N}{2}} \exp \bigg(\frac{iN_1 \lambda_\mathrm{nc} \tilde{\bar Q}}{1 - 2i \tilde{\bar Q}} \bigg), \nonumber\] we find with \(\lambda_\mathrm{nc} = \frac{\bar f^\top m m^\top \bar f}{\bar f^\top \Sigma \bar f}\) \[\begin{align} \varphi (\bar f| X) & = \int \frac{d\tilde{\bar Q} d\tilde{Q}}{2\pi} \big(1 - 2i \tilde{\bar Q}\big)^{-\frac{N}{2}} \exp \bigg[ \frac{i}{1 - 2i \tilde{\bar Q}} \frac{N_1 \tilde{\bar Q}}{\bar f^\top \Sigma \bar f} \bar f^\top m m^\top \bar f + i \tilde{Q} \left( -\tilde{\bar Q} + \frac{i}{2 N_1 \lambda_1} \bar f^\top \Sigma \bar f \right) \bigg]. \end{align}\] Noting that the second term in the exponent gives precise meaning to \(\tilde{\bar Q}\), as \(\int d\tilde{Q} e^{i \tilde{Q} \left( -\tilde{\bar Q} + \frac{i}{2 N_1 \lambda_1} \bar f^\top \Sigma \bar f \right)} = \delta (\tilde{\bar Q} - \frac{i}{2 N_1 \lambda_1} \bar f^\top \Sigma \bar f)\), we can cancel \(\frac{N_1 \tilde{\bar Q}}{\bar f^\top \Sigma \bar f} = \frac{i}{2 \lambda_1}\) in the first term of the exponent - this leaves the desired quadratic function of \(\bar f\), and with the substitutions \(- 2i \tilde{\bar Q} = \bar Q\) and \(\tilde{Q} / N = Q\) \[\begin{align} \varphi (\bar f| X) & = \int \frac{d\bar Q d\tilde{Q}}{2\pi} \big(1 - 2i \tilde{\bar Q}\big)^{-\frac{N}{2}} \exp \bigg[ - \frac{1}{2 \lambda_1}\frac{1}{1 - 2i \tilde{\bar Q}} \bar f^\top m m^\top \bar f - i \tilde{Q} \tilde{\bar Q} - \frac{\tilde{Q}}{2 N_1 \lambda_1} \bar f^\top \Sigma \bar f \bigg] \\ & \propto \int d\bar Q d\tilde{Q} \big(1 + \bar Q\big)^{-\frac{N}{2}} e^{\frac{N_1}{2} Q \bar Q }\, \exp \bigg[ - \frac{1}{2} \bar f^\top \bigg( Q\, \Theta(C) + \big( \frac{1}{1 + \bar Q} - Q \big) \frac{m m^\top}{\lambda_1} \bigg) \bar f \bigg] \\ & \propto \int d\bar Q d\tilde{Q} \exp \bigg[ \frac{N_1}{2} Q \bar Q - \frac{N_1}{2} \log (1 + \bar Q) - \frac{1}{2} \bar f^\top \big( \mathbb{1}\beta^{-1} + K^R_{\mathrm{nc}}(Q,\bar Q)\big) \bar f \bigg]. \label{eq:app95charfunc95prior951hl95noncentral} \end{align}\tag{54}\] In the last line we defined the renormalized kernel for the non-central EWA, \[\begin{align} K^R_{\mathrm{nc}} (Q,\bar Q) &= K^R (Q) + \frac{1}{\lambda_1} \bigg( \frac{1}{1 + \bar Q} - Q \bigg) m m^\top \\ &= Q\, \Theta(C) + \frac{1}{\lambda_1} \bigg( \frac{1}{1 + \bar Q} - Q \bigg) m m^\top. \label{eq:renormalized95noncentral951hl95kernel} \end{align}\tag{55}\] Finally, plugging into the expression for the partition function of the posterior Eq. 11 and performing the Gaussian integration, we find \(Z_\mathrm{nc} \propto \int d\bar Q d\tilde{Q}\, e^{- \frac{N_1}{2} S[Q,\bar Q]}\) with the effective action \[S_\mathrm{nc} \big[Q,\bar{{Q}}\big] = -Q\bar{Q} +\log\left(1+ \bar Q\right) +\frac{{\alpha}}{P} \log\det\left[ \mathbb{1}\beta^{-1} + K_{\mathrm{nc}}^{R} (Q,\bar Q )\right] +\frac{{\alpha}}{P} y^{\top} \left(\mathbb{1}\beta^{-1} + K_{\mathrm{nc}}^{R} (Q,\bar Q )\right)^{-1}y. \label{eq:1hl-raw-noncentral-action}\tag{56}\] This recovers the one hidden layer action for non-zero mean activations in Suppl. Sect. IV of [40], here obtained via the route of a non-central EWA. By analyzing this expression further we here show explicitly how the difference to the central EWA is small, as argued generally in Section 4.1.

Using the Sherman–Morrison formula \[\left(A+uv^{\intercal}\right)^{-1}=A^{-1}-\frac{A^{-1}uv^{\intercal}A^{-1}}{1+v^{\intercal}A^{-1}u}\] and the Matrix determinant lemma \[\det\left(A+uv^{\intercal}\right)=\left(1+v^{\intercal}A^{-1}u\right)\det\left(A\right),\] 56 can be expressed in the more lengthy but better interpretable form written only in terms of the kernel \(K^R_\beta(Q)\) appearing in the cental EWA: \[\begin{align} S_\mathrm{nc} \left[Q,\bar{Q}\right] =& -Q\bar{Q} + \log\left(1+\bar{Q}\right) +\frac{{\alpha}}{P} \mathop{}\!\mathrm{Tr}\log \left[\mathbb{1}\beta^{-1} + K^R\right] +\frac{{\alpha}}{P} y^{\top}\left(\mathbb{1}\beta^{-1} + K^R\right)^{-1} y \nonumber \\ & +\frac{{\alpha}}{P} \log \left[ 1 - \frac{{1}}{\lambda_{1}} \left(Q-\frac{{1}}{1+\bar{Q}}\right) m^{\top}\left(\mathbb{1}\beta^{-1} + K^R\right)^{-1}m \right] \label{eq:noncentral951hl95action95expanded}\\ & +\frac{{\alpha}}{P} y^{\intercal}\left[ \frac{{1}}{\lambda_{1}} \left(Q-\frac{{1}}{1+\bar{Q}}\right) \frac{\left(\mathbb{1}\beta^{-1} + K^R\right)^{-1} mm^{\top} \left(\mathbb{1}\beta^{-1} + K^R\right)^{-1}}{1 - \frac{{1}}{\lambda_{1}} \left(Q-\frac{1}{1 + \bar{Q}}\right) m^{\top}\left(\mathbb{1}\beta^{-1} + K^R\right)^{-1} m} \right]y \nonumber . \end{align}\tag{57}\] Here the first line is equivalent to the action Eq. 35 obtained from the central EWA, as seen by plugging in the explicit saddle-point relation \(\frac{1}{Q^{\ast,c}} - 1 = \bar{Q}^{\ast,c}\) which follows directly from \(\partial S / \partial \bar{Q} \overset{!}{=} 0\) when setting the second and third line to zero. The second and third line thus represent the difference between central and non-central EWA. Note that at the saddle-point of the central EWA, the quantity \(\Delta_Q \coloneq Q-\frac{{1}}{1+\bar{Q}} \overset{\ast,c}{=} 0\) causing the second and third lines to vanish. While the saddle point of the non-central action Eq. 57 can give rise to \(\Delta_Q \neq 0\) at finite \(N,P\), we find that asymptotically \(\Delta_{Q^\ast} \overset{N,P\to\infty}{\to} 0\).

Zero-temperature limit↩︎

For ease of exposition, we now focus on the zero-temperature limit where the structure of this action becomes particularly clear. Setting \(\beta \to 0\), we can define the task-dependent (but not \(Q, \bar{Q}\) dependent) scalar overlaps \[\begin{align} M_{yy} &= \frac{1}{P} y^{\intercal} \Theta(C)^{-1} y \tag{58}\\ M_{m y} &= \frac{1}{\sqrt{P}} m^{\intercal} \Theta(C)^{-1} y \\ M_{mm} &= m^{\intercal} \Theta(C)^{-1} m \\ \Gamma_K &= \frac{1}{P} \log [\det \Theta(C)]. \tag{59} \end{align}\] With the short-hand \(\Delta_{Q} = Q-\frac{1}{1 + \bar Q}\), the action becomes \[\begin{align} S[Q,\bar Q] = &- Q \bar Q + \log\left(1+\bar Q\right) + \alpha \log(Q/\lambda_1) + \alpha \Gamma_K + \alpha\frac{\lambda_1}{Q} M_{yy} \nonumber \\ &+ \frac{\alpha}{P} \log \left( 1 - \frac{\Delta_Q}{Q} M_{m m} \right) + \alpha \lambda_1 \frac{\Delta_Q}{Q \left(Q - \Delta_Q \;M_{m m} \right)} M_{m y}^2. \label{eq:noncental951hl95action95zerotemp} \end{align}\tag{60}\] Again, the first line corresponds to the central EWA result. Note that the saddle-point of the zero temperature, central EWA action only depends on \(\alpha\) and the fixed task-to-kernel overlap \(M_{yy}\), permitting a simple understanding of the the saddle-point solution \(Q^\ast\), see App. 14.

Scaling behavior↩︎

The starting observation to analyze the relative scaling and contribution of the \(M_{yy},M_{\mu y},M_{\mu \mu}\) overlaps is that the NNGP kernel \(\Theta(C)\) corresponds to to a zero-mean activation kernel with rank-1 spike \(\Theta(C) = \Sigma + m m^\top.\) Since \(m^\mu =O(1)\), while the BBP transition happens at the \(O\big(1/\sqrt{P}\big)\) scale, the spike always creates a strong outlier Eigenvalue \(\lambda_0=O(P)\) as long as \(\Sigma\) does not inherit an equally strong low-rank structure from the data distribution. Diagonalizing the kernel into the Eigenvalue-Eigenvector pairs \(\{\lambda_i, v_i\}_{i=1..P}\) we can define the projections of \(m\) and \(y\) on the Eigenspaces, \(\hat{m}_i = m^\top v_i\) and \(\hat{y}_i = y^\top v_i\). When inverting the kernel, the Eigenvectors stay the same and the Eigenvalues are simply inverted, giving the pairs \(\{\lambda_i^{-1}, v_i\}_{i=1..P}\).

Therefore, we can understand the the scaling of \(M_{yy},M_{\mu y},M_{\mu \mu}\) through the behavior of the kernel spectrum and the alignment of the Eigenspaces with \(m,y\). In particular, \(v_0\) is typically highly localized to \(m\) and the outlier eigenvalue \(\lambda_0^{-1}\) becomes the smallest eigenvalue of the inverse kernel, which instead is dominated by the tail of the spectrum. The exact scaling indeed still depends non-trivially on the data-dependent decay of the spectrum and the relative overlaps of \(m,y\) with the corresponding Eigenspaces. To avoid a lengthy discussion we rely on the general argument of asymptotic equivalence to the central EWA presented in Section 4.1, and here restrict ourselves to present numerical results on the behavior of \(M_{yy},M_{\mu y},M_{\mu \mu}\) and \(\Delta_{Q^\ast}\) in Appendix 9.3 and a few observations:

• For real data distributions such as the images of CIFAR-10, the dimensionality of the ground-truth data manifold is fixed as the number of samples \(P\) taken from the distribution increases. A well known consequence is that the Eigenvalues of the kernel at fixed spectral index \(i\) grow as \(\lambda_i \sim P\), while in relative terms the new eigenvalues added in the tail of the spectrum as \(P\) increases, assuming the modes at \(i=O(P)\) still approximate the population spectrum, typically follow a power-law decay \(\lambda_{O(P)} /P \sim P^{-\gamma}\) with kernel- and data-dependent exponent \(\gamma > 1\). Such a power law decay holds for the population spectrum of finite-smoothness kernels such as the ReLU NNGP [82], [83], while for analytic kernels the decay may be faster. Therefore, the tail comes to dominate the inverse spectrum more and more, e.g. here with the mean \(\frac{1}{P}\sum_i\lambda_i^{-1} \sim P^{\gamma -1}\).

• Correspondingly, if the label vector \(y\) projects uniformly on the Eigenspaces of the kernel, also \(M_{yy}\) grows as \(\sim P^{\gamma - 1}\). For easier tasks where \(y\) is largely localized to the upper bulk of the spectrum, \(M_{yy}=O(1)\).

• The logarithmic term \(\frac{\alpha}{P} \log \left( 1 - \frac{\Delta_Q}{Q} M_{m m} \right)\) trivially vanishes for \(P\to \infty\) as long as \(\frac{\Delta_Q}{Q} M_{m m} \ll 1\) to avoid the divergence. This behavior is encouraged both since \(M_{m m}\overset{P\to\infty}{\to}0\) if \(v_0\) is mostly localized to \(m\), and since the saddle-point can be thought of as a perturbed version of the central EWA saddle located at \(\Delta_{Q^{\ast,c}}=0\).

• \(M^2_{my}\) which controls the last term in Eq. 60 is determined by the overlap of the projections of \(m\) and \(y\) on the Eigenspaces. The contributions of each \(\hat{m}_i \lambda_i^{-1} \hat{y}_i\) can be positive or negative and therefore partly cancel. If \(m||y\), then \(M_{my} \propto \frac{1}{\sqrt{P}} M_{mm}\) vanishes at a faster rate than \(M_{mm}\). A fast decay exponent \(\gamma\) of the tail spectrum can lead to slowly decaying or even constant \(M_{my}\), however the same mechanism increases also \(M_{yy}\), such that the hierarchy between the two terms remains.

9.3 Numerical evidence of the equivalence of central and non-central approaches↩︎

The behaviors discussed above are shown numerically in Fig. 9 for the CIFAR-10 \(\{0,1\}\) and Gaussian data tasks used in the main text. For MNIST, and synthetic datasets with class imbalance or with labels \(y\) correlated to the mean vector \(m\), we have found qualitatively similar results (not shown). In all cases, and as following from the general arguments presented in the main text Section 4.1, the quantity \(\Delta_Q\) controlling any differences between central and non-central theories vanishes with increasing system size \(P,N\). For the synthetic Gaussian task this happens as \(\Delta_Q\sim P^{-1}\). The second column of Fig. 9 shows this is due to similarly vanishing \(|M_{my}|\).

As seen in the third column panels, the majority contribution to \(M_{mm}\) is due to the outlier eigenvalue. This shows that also in the fixed data dimension case of CIFAR, where not only the mean-associated outlier but all eigenvalues grow proportionally with \(P\), the outlier eigenvector \(v_0\) is largely localized and \(m\) has non-vanishing overlaps only with a few top eigenvectors. \(M_{my}\) also shows a significant contribution from the outlier eigenvalue, and while the contributions of remaining directions are of significant size they have (for these choices of task and kernel) almost random signs, causing these contributions to cancel. To visualize the size of this cancellation, the green lines in the third column were normalized by the total size of unsigned contributions. Finally, the contributions to \(M_{yy}\) are more evenly spread across the full spectrum. This is task dependent but expected of a regime where adding data samples increases the accuracy of the kernel predictor.

\(M_{yy}\) can be growing with \(P\), as in the case of CIFAR where numerically \(M_{yy}\sim P^{0.3}\). This growth is due to the increasing weight of the tail eigenvalues; the last 30% of eigenvalues for each \(P\) contribute about half of the size of \(M_{yy}\) for CIFAR (orange lines in third column), while the mean value of the inverse spectrum increases with \(P\) (inset of fourth column).

Lastly, as expected the inverse NNGP spectrum of the Gaussian data (bottom right panel) is composed of three easily interpretable parts: The outlier eigenvalue due to the non-zero mean scaling with \(P\), and two bulks of eigenvalues. These arise from the fact that \(C=XX^\top\) is Wishart and not of full rank since \(P=2N\), creating two bulks of \(N\) nonzero and \(N\) zero eigenvalues. Trivially the linear label vector \(y=w^\ast X\) can only have overlap with the nonzero part of the spectrum, which explains the behavior of the \(M_{yy}\) and \(M_{my}\) contributions in the third column. The ReLU NNGP kernel function \(\Theta\), having an infinite-dimensional associated reproducing kernel Hilbert space (RKHS), ensures that \(\Theta(C)\) is full-rank such that also the second bulk of eigenvalues is non-zero, independently of small regularization (compare separation of gray dashed line and inverse spectrum).

9.4 Derivation of Eq. 28 : \(t=\mathop{}\!\mathrm{Tr}(\Sigma K_\beta)\) controls self-averaging of \(\bar f^\top\Sigma\bar f\)↩︎

This appendix provides a short derivation of the quantity \(t=\mathop{}\!\mathrm{Tr}(\Sigma K_\beta)\), which is used in Sec. 4.1 to show that under the \(\bar f\) measure the typical noncentrality parameter \(\lambda_\mathrm{nc}=(\bar f^\top m)^2/(\bar f^\top\Sigma\bar f)\) is \(O(1/P)\) and negligible.

In the noncentral EWA one encounters \(\bar f\) integrals of the form \[\int d\bar f\; \exp\!\Big(-\tfrac12 \bar f^\top K_\beta^{-1}\,\bar f + i y^\top \bar f\Big)\,(\cdots), \qquad K_\beta=(\beta^{-1}I+Q\,\Theta)^{-1}, \label{eq:appteff95shifted95gaussian}\tag{61}\] with \(\Theta=\Sigma+mm^\top\) and \(Q>0\). Completing the square gives \[-\tfrac12 \bar f^\top K_\beta^{-1} \bar f+i y^\top \bar f = -\tfrac12(\bar f-\mu)^\top K_\beta^{-1} (\bar f-\mu)+\tfrac12\,y^\top K_\beta \,y, \qquad \mu\coloneq i K_\beta \, y.\] Thus the \(\bar f\)-dependence in Eq. 61 is that of a complex-mean Gaussian with covariance \(K_\beta\). In particular, fluctuations of quadratic forms such as \(\bar f^\top\Sigma\bar f\) are then governed by \(K_\beta\), while the linear term only adds the deterministic shift \(\mu\). Write \(\bar f=\mu+g\) with \(g\sim\mathcal{N}(0,K_\beta)\). Then \[D=\bar f^\top\Sigma\bar f = g^\top\Sigma g + 2\mu^\top\Sigma g + \mu^\top\Sigma\mu \label{eq:appteff95D95decomp}\tag{62}\] and mean and variance of the pure fluctuation term are \[\begin{align} \mathbb{E}[g^\top\Sigma g] &= \mathop{}\!\mathrm{Tr}(\Sigma K_\beta) \equiv t, \\ \mathrm{Var}[g^\top\Sigma g] &= 2\,\mathop{}\!\mathrm{Tr}\!\big((\Sigma K_\beta)^2\big). \end{align}\] Furthermore, \(\Sigma\) and \(K_\beta^{-1}=\beta^{-1}I+Q\,(\Sigma + m m^\top)\) are positive semi-definite and the eigenvalues of \(\Sigma K_\beta\) are \(\in(0,Q^{-1}]\), so that \[\mathop{}\!\mathrm{Tr}\!\big((\Sigma K_\beta)^2\big)\le Q^{2}\mathop{}\!\mathrm{Tr}(\Sigma K_\beta) = Q^2 t \qquad\Rightarrow\qquad \frac{\sqrt{\mathrm{Var}[g^\top\Sigma g]}}{\mathbb{E}[g^\top\Sigma g]} \le \sqrt{\frac{2}{t}}Q.\] Hence whenever \(t=\mathop{}\!\mathrm{Tr}(\Sigma K_\beta)\) diverges, \(g^\top\Sigma g\) is self-averaging. In the main text, it is argued that \(t=O(P)\) and therefore relative fluctuations are \(O(P^{-1/2})\).

The mean-shift related terms in Eq. 62 do not change this behavior: \(2\mu^\top\Sigma g\) is linear in \(g\) and has variance \(4\,\mu^\top\Sigma K_\beta\Sigma\mu\), while \(\mu^\top\Sigma\mu\) is deterministic. For the \(\|y\|^2=O(P)\) scaling relevant here, these contributions are at most \(O(P)\) even if \(y\) correlates significantly with the mean direction \(m\), and therefore do not affect the fact that \(D\) is extensive and concentrates around its mean on relative scale \(P^{-1/2}\).

10 Additional details on the Large Deviation Analysis↩︎

We here discuss more in depth some aspects on the large deviation analysis performed in Sec. 4.2. Recall that the variables under investigation are the \(q_\ell\) variables, defined as \[\label{eq:q95ell} q_\ell = \frac{\bar f^\top \Theta^{L-\ell}(K_\mathrm{E}^{(\ell)}) \bar f}{\bar f^\top \Theta^{L-\ell+1}(K_\mathrm{E}^{(\ell-1)}) \bar f}\tag{63}\] and viewed as sequences over the corresponding \(N_\ell\). We keep a generic \(N_\ell\), but all the plots will show for simplicity results obtained by sampling networks with the same number of neurons across layers: \(N_\ell=N \, \forall \ell\). Here we distinguish explicitly between central and non-central EWA, and we will provide additional numerical evidence on the claims made in the main text.
We first show that the normalized quantity in Eq. 63 , introduced in the main text in the context of the central EWA, is actually left unchanged even for the non-central case. Indeed, as shown in Sec. 9.1, the non-central EWA implies that \(\tilde{Q}_\ell \sim \chi_{\text{nc}}^2(N_\ell, N_\ell \lambda_{\text{nc},\ell})\). As a consequence, the normalized variable: \[\begin{align} \tilde{q}_\ell &\coloneq \frac{\tilde{Q}_\ell}{\mathbb{E}[\tilde{Q}_\ell]} = \frac{\tilde{Q}_\ell}{N_\ell(1+\lambda_{\text{nc},\ell})} = q_\ell. \end{align}\] Where the last equality follows from Eq. 53 . In particular, we can use the same sampling algorithm for both the central and non central case, the only difference being the specific function implemented by the NNGP kernel function \(\Theta\). Naturally, the central and non-central EWA will generally lead to a different shape of the theoretical rate function, as will be now described, but the way in which \(q_\ell\) and \(\tilde{q}_\ell\) are calculated is the same, even if their distributions are different. In order to derive the layer-wise rate function for the non-central EWA, let us consider again the fact that \(\tilde{Q}_\ell^{(N_\ell)} \sim \chi^2_{\text{nc}}(N_\ell, N_\ell \lambda_\ell)\). The moment generating function of the non-central chi-squared distribution is: \[M_{\tilde{Q}_\ell^{(N_\ell)}}(t) = (1-2t)^{-N_\ell/2} \exp \left\{ \frac{N_\ell \lambda_{\text{nc},\ell} t}{1-2t} \right\}, \qquad t<1/2.\] It follows that the moment generating function of the normalized \(\tilde{q}_\ell\) are: \[M_{\tilde{q}_\ell^{(N_\ell)}}(t) = \left(1-\frac{2t}{N_\ell(1+\lambda_{\text{nc},\ell})}\right)^{-N_\ell/2} \exp \left\{ \frac{N_\ell \lambda_{\text{nc},\ell} t}{N_\ell(1+\lambda_{\text{nc},\ell})-2t} \right\}, \qquad t<N_\ell(1+\lambda_{\text{nc},\ell})/2.\] Therefore, the scaled cumulant generating functions are: \[\begin{align} \Lambda_\ell(t) &= \lim_{N_\ell\to\infty} \frac{1}{a_{N_\ell}} \ln M_{\tilde{q}_\ell^{(N_\ell)}}(a_{N_\ell}t) \\ &= \frac{\lambda_{\text{nc},\ell} t}{1+\lambda_{\text{nc},\ell}-2t} - \frac{1}{2} \ln \left( 1-\frac{2t}{1+\lambda_{\text{nc},\ell}} \right), \qquad t<(1+\lambda_{\text{nc},\ell})/2. \end{align}\] where the scales for these LDP turned out to be \(a_{N_\ell}=N_\ell\). Finally, the layer by layer rate functions for the non-central EWA are obtained via Legendre–Fenchel tranform: \[\mathcal{I}_\ell (x) = \sup_{t<\frac{1+\lambda_{\text{nc},\ell}}{2}} \left\{ tx - \frac{\lambda_{\text{nc},\ell} t}{1+\lambda_{\text{nc},\ell}-2t} + \frac{1}{2} \ln \left( 1-\frac{2t}{1+\lambda_{\text{nc},\ell}} \right) \right\}.\] And this expression can be easily computed numerically. Notice that this time, differently than in the central case, the rate function is in principle different for different layers, but only through the non centrality parameter \(\lambda_{\text{nc},\ell}\), recovering the central result for \(\lambda_{\text{nc},\ell}=0\). This is due to the fact that, differently than in the central case, the \(\tilde{q}_\ell\) variable are not decoupled across layers, since they conditionally depend on the previous layer kernel through the non centrality parameter, which we recall it is given by: \[\begin{align} \lambda_{\mathrm{nc},\ell} = \frac{\bar{f}^\top m_{(\ell)} m_{(\ell)}^\top \bar{f}}{\bar{f}^\top \Sigma_{(\ell)} \bar{f}}, \end{align}\] where \(m_{(\ell)}\) and \(\Sigma_{\ell}\) are calculated using \(K_\mathrm{E}^{(\ell-1)}\), see Sec. 9.1. In our numerical experiments, we choose to control not the \(\lambda_{\textrm{nc}}\) parameter but the overlap between \(\bar f^\top\) and the last layer \(m_{(L)}\), since this is the quantity that quantifies the contributions of non-central effects at the posterior level, so in the final integration over \(d\bar f\). For this reason, we need to modify the sampling algorithm used in the main text for the central case in order to get proper samples from each of the conditional distributions \(\tilde{q}_\ell \mid K_{\mathrm{E}}^{(l-1)}\). On a practical level, the difference between this modified sampling algorithm, that we call Conditional Sampling Algorithm, and the simpler one used to sample from the central joint distribution is that in this case the denominator of Eq. 63 is kept fixed across samples. It follows in particular that there is no difference between the two sampling procedures for \(1\)-hidden layer. The scheme of the algorithm is summarized in Alg. 10.

As described in the main text, in the case of the central EWA independent samples from the joint distribution of \((q_1, \dots, q_L)\) are needed, in order to get proper samples of random variable \(\mathcal{Q}=\prod_\ell q_\ell\). Even though the predictive power of the EWA only require that the variable \(\mathcal{Q}\) satisfies the LDP, we also provide numerical results concerning the LDP at the level of each individual layer even in the central case. Additionally, in order to numerically check that the \(q_\ell\) variables are indeed decoupled across layers, as predicted by the central EWA, we use the conditional algorithm also for the central case: if the corresponding empirical rate functions are in agreement with the theoretical ones predicted by the EWA it means that the \(q_\ell\) are indeed decoupled. Fig. 11 shows numerical evidence on how the EWA is able to catch the statistical properties of the \(q_\ell\) variable even at the level of each individual layer, especially in the Erf case (first row). We systematically observe some asymptotic deviation for ReLU networks (second row), suggesting that the deviation is then absorbed at the product level, see Fig. 2 and the following Sec. 10.1 for an explanation using a toy model with independent variables.

It is worth mentioning how the result of the sampling procedure is stable across different realization of the \(\bar f\) vectors: From Eq. 63 it follows that the normalization of the \(\bar f\) is not relevant, and different realization of the \(\bar f\) vectors sampled uniformly from the unit \(P\)-dimensional sphere leads to the same conclusions, as expected. See Fig. 12 for an example.

Fig. 13 shows instead how the non-central EWA is conceptually a more principled ansatz at the prior level, in the case of a shallow architecture for which the effective action was derived in Sec. 9.2. The plot shows how the theoretical rate function of the non-central EWA is able to describe the empirical sampled points, especially for large values of the overlap parameters when the difference between central and non-central EWA becomes more evident. The reason why these differences become irrelevant at the level of the prior and posterior output distribution in the proportional regime is explained in Sec. 4.1.

10.1 Toy model for deviations from the theoretical product rate function↩︎

To understand how deviations in the width of layer-wise rate functions influence the product rate function, consider a simple toy model. Assume that the variables \(x_\ell >0\) are independent and admit rate functions of the same speed of the form \[I_\ell(x) = a_\ell I_x(x),\] where \(I_x\) is a common base rate function and the prefactors \(a_\ell\) represent quenched disorder around the homogeneous value \(a_\ell=1\). For the product \(y=\prod_{\ell=1}^L x_\ell\) the contraction principle gives \[I_y(y;\{a_\ell\})=\inf_{\prod_{\ell=1}^L x_\ell = y}\sum_{\ell=1}^L a_\ell I_x(x_\ell).\] Expanding around the typical point \(x_\ast=\mathrm{argmin\,I_x}\) such that \(I_x(x)\simeq \frac{\kappa}{2}(x-x_\ast)^2,\) and writing \(x_\ell=x_\ast+\delta_\ell\), the product constraint expands to leading order as \(x_\ast^{L-1}\sum_{\ell=1}^L \delta_\ell \simeq y - y_\ast,\) with \(y_\ast=x_\ast^L\). The rate function then reduces to the quadratic minimization problem \[I_y(y;\{a_\ell\}) \simeq \inf_{\sum_\ell \delta_\ell = \frac{y-x_\ast^L}{x_\ast^{L-1}}}\frac{\kappa}{2}\sum_{\ell=1}^L a_\ell \delta_\ell^2,\] whose solution is \[I_y(y;\{a_\ell\}) \simeq \frac{\kappa}{2}\,\frac{(y-x_\ast^L)^2}{x_\ast^{2L-2}\sum_{\ell=1}^L a_\ell^{-1}} = \bigg(\frac{1}{L} \sum_{\ell=1}^L a_\ell^{-1}\bigg)^{-1} I_y(y;\{a_\ell =1\}) .\] Thus the disorder enters only through the harmonic mean \(\frac{1}{L}\sum_\ell a_\ell^{-1}\).

This model with independent \(x_\ell\) immediately explains two qualitative features. First, the effect of the width-disorder is self-averaging and not cumulative. The homogeneous approximation therefore does not deteriorate in the product but can often be better than at individual layers. Second, the harmonic mean makes the rate function less sensitive to large values \(a_\ell>1\) than to small values \(a_\ell<1\), since increasingly large values of any given \(a_\ell\) decrease \(a_\ell^{-1}\) only weakly, whereas decreasing \(a_\ell\) below unity enhances \(a_\ell^{-1}\) more than proportionally. This corresponds to the dominance of the softest modes in the tails of the distribution of the product, while stiffer modes contribute less.

11 Ease and difficulty of sampling the overparametrized DNN posterior↩︎

Generally, sampling from the posterior of a DNN is a formidable task ue to its very high dimensionality and multi-modality. From the outset, it can only become feasible since we are interested solely in low-dimensional summary statistics such as the generalization error, or at most the marginals of the output posterior (\(P\) scalar distributions), instead of the full joint distribution in \(N\times N\) parameter space. Difficulties can especially be expected at low temperature and high network load \(\alpha = P/N\). In such cases, due to the non-convex loss landscape in the parameter space, Monte Carlo sampling could realistically exhibit several systematic biases in the exploration of the configuration space, breaking convergence to the true posterior on any feasible time scale and trustworthiness of numerical outcomes. There are multiple reasons to expect a breakdown of Monte Carlo sampling: for example, networks can get stuck in individual branches of the non-convex loss landscape and require exponential time to escape [65], or the simulation can slow down due to poor gradient propagation across layers such that some parameter groups are sampled on much slower time scales than others [84]. Nevertheless Monte Carlo practitioners have different strategies to numerically probe the stability of their simulations. Below, we present the algorithms we employed for sampling, together with the statistical tools we implemented to assess the quality of our simulations. Discussing examples representative of our results, we show that at our working scale Bayesian experiments involving non-linear DNNs are feasible.

11.1 Additional MCMC samplers↩︎

A possible approach to probe the convergence of Monte Carlo simulations is to compare outcomes from different algorithms. For this reason, in hard regimes where one simultaneously needs large depth and training set size we employed several different algorithms. We chose algorithms that do not share the same paradigm for exploring the network configuration space. In particular, we implemented a pure gradient-based Monte Carlo sampler (Langevin Monte Carlo – LMC), a pure energy-based Monte Carlo sampler (the preconditioned Crank–Nicolson MCMC algorithm – pCN), a mixed gradient-based and energy-based Monte Carlo algorithm (the Metropolis Adjusted Langevin Algorithm – MALA), and a momentum-driven Monte Carlo algorithm (the Hamiltonian Monte Carlo algorithm with No-U-Turn Sampler – NUTS HMC). In the manuscript, Sec. 6.2, we already discussed the LMC implementation, since we mainly used this algorithm for testing the EWA in different learning scenarios (see the section below for details on how we assessed the sampling robustness of this algorithm). Here we provide a brief description of the routines we used to implement the other algorithms:

• Metropolis Adjusted Langevin Algorithm – MALA
  The MALA algorithm is an exact Bayesian sampler built on the discretized Langevin equation. While more expensive, its main advantage compared to LMC is that it samples from the true posterior, avoiding additional systematics arising from the finite learning rate. Since the standard discretized Langevin transition \[\label{eq:transitions} \theta \rightarrow \theta' = \theta - \epsilon \partial_\theta\left[ \mathcal{L}_{\mathrm{reg.}}\right]_{\theta} + \sqrt{2\epsilon T}\eta\tag{64}\] is not symmetric, i.e. \(P_\beta(\theta\to\theta')\ne P_\beta(\theta'\to\theta)\), it is not enough to correct the transition \(\theta\to\theta'\) with a Metropolis step to ensure that detailed balance is satisfied. On the contrary, one has to take into account the full transition probability between two states \(\theta\) and \(\theta'\), which basically amounts to accepting the new state \(\theta'\) as a proposal from state \(\theta\) with probability \[\begin{align} \label{eq:acc95prob} p_\beta(\theta, \theta') = \min \left( 1, \frac{P_\beta(\theta')P_\beta(\theta'\to\theta)}{P_\beta(\theta)P_\beta(\theta\to\theta')} \right)\, , \end{align}\tag{65}\] where the transition amplitudes can be easily obtained from Eq. 64 , namely: \[\label{eq:MALA95transitions} P_\beta(\theta\to\theta') = P_\beta(\eta = [\theta' - \theta + \epsilon \partial_\theta \mathcal{L}_{\mathrm{reg.}}]/\sqrt{2\epsilon T} ) \propto \exp{\left\{-\frac{1}{4\epsilon T}\left( \theta' - \theta + \epsilon \partial_\theta \mathcal{L}_{\mathrm{reg.}} \right)^2\right\}}\, .\tag{66}\] So, MALA simply amounts to iterating the gradient-based proposal in Eq. 64 and the energy-based acceptance step in Eq. 65 . The price one pays to avoid finite learning rate effects is not only the computation of the loss at each step, but also the computation of the transition amplitudes. We note also from Eqs. 65 and 66 that the acceptance rate depends on the choice of learning rate, requiring in principle extra preliminary simulations to fix the value of \(\epsilon\) for each set of parameters in order to obtain a fixed target acceptance probability. In our simulations we fixed \(\epsilon=10^{-3}\), which allowed us to keep the average acceptance probability in the range \(p_\beta\approx 0.8\) for \(L=5\) and \(p_\beta\approx 0.4\) for \(L=10\).

• Preconditioned Crank–Nicolson MCMC algorithm – pCN
  The preconditioned Crank–Nicolson is a Monte Carlo sampling algorithm with the special feature of preserving a good acceptance rate even in the high-dimensional limit (which is not the case for MALA, for example, where \(p_\beta\to0\) as \(N_\ell\to\infty\)). The pCN proposal is gradient-free, namely: \[\theta\to\theta'=\sqrt{1-\phi^2}\theta + \frac{\phi}{\sqrt{\lambda}}\xi\, , \qquad \xi\sim\mathcal{N}(0,\mathbb{1})\, ,\] where \(\phi\in [0,1]\) is a free parameter to be optimized to ensure a reliable acceptance rate. The new proposal \(\theta'\) is accepted with probability \[\label{eq:acc95step95pcn} p_\beta(\theta,\theta')=\min[1, \exp(-\beta(\mathcal{L}(\theta')-\mathcal{L}(\theta)))] \, .\tag{67}\] Note that, unlike MALA, the acceptance step of pCN is computed only on the likelihood and not on the full regularized loss. It is also worth mentioning that the pCN proposals always remain on the shell of the prior norm, rather than drifting away from it, diffusing through the acceptance step in Eq. 67 . In this work, we fixed \(\phi=0.002\), which yields an average acceptance probability in the range from \(p_\beta \approx 0.48\) (for \(L=5\)) to \(p_\beta\approx 0.27\) (for \(L=10\)).

• No-U-Turn Hamiltonian Monte Carlo – NUTS HMC
  The NUTS HMC algorithm is a state-of-the-art sampler which is particularly suitable for Bayesian models. While it is well known that the vanilla HMC algorithm is a sampler with fast mixing capabilities due to momentum-driven phase-space exploration (resulting in short decorrelation timescales), it is also well known that HMC requires preliminary runs to set the hyperparameters of the simulation, to which the performance of the algorithm is very sensitive. The most crucial hyperparameters of the simulation are the integration step size and the trajectory length. In particular, a wrong choice of the latter is chancy for the simulation, as it can yield trajectories which fold back onto the same track. NUTS HMC builds on top of the HMC algorithm and solves the problem of fixing the trajectory length. Roughly speaking, the algorithm tries different trajectories forward and backward in time, increasing the trajectory length until the system is far from loops back to a previous position in the trajectory (in jargon, a U-turn), and samples one trajectory among those generated this way to ensure detailed balance. A full description of the sampling strategy, which is beyond the scope of this manuscript, is presented in Ref. [68]. In this work we used the NUTS HMC sampler provided by the NumPyro Python package [85], which also provides optimal fine-tuning of the learning rate. It can be shown that this algorithm performs at least as well as optimally fine-tuned HMC [68]. Before starting the sampling, we fixed a pre-calibration phase of \(500\) steps: during this warm-up phase, the system was first brought to thermalization in \(O(10)\) steps, and then the remaining steps were used to fix the optimal trajectory length and integration step size to achieve an optimal acceptance rate. Typical values of the step size \(\epsilon\) are in the range \(10^{-4}<\epsilon<10^{-3}\), while for the trajectory length \(\tau\) we obtained an optimal \(\tau \approx O(10^3)\), which yielded typical average acceptance probabilities \(0.7<p_\beta<0.95\).

11.2 Methods used to assess sampling convergence↩︎

The estimation of the autocorrelation time \(\tau_{\text{int}}\) [64] is a well-known tool for the computation of statistical errors associated with Monte Carlo measurements. In this work we used the blocking method, which gives an indirect measure of the autocorrelation time by means of the statistical estimation of the standard deviation of the mean. It operates by partitioning the data into blocks and averaging the values within each block, thereby creating new samples of reduced length and decreased autocorrelations. As the block size grows, the error estimate obtained from the naive standard error of the mean (i.e. neglecting autocorrelation effects) on the blocked samples becomes increasingly accurate, eventually converging to a plateau once the blocks are effectively uncorrelated. The plateau value \(\delta\) corresponds to the actual statistical error associated with the mean, and the integrated autocorrelation time can be computed using the relation \(\delta^2 = \sigma_0^2 2\tau_{\text{int}}/N_{\text{samples}}\), where \(\sigma_0\) is the standard deviation of the original samples. Note that from the previous formula one can also obtain the number of independent samples \(N_{\text{ind. samples}}=N_{\text{samples}}/(2\tau_{\text{int}})\). In Fig. 14 (b) we show a representative example of the estimation of the statistical error using the blocking method in the case of a MLP with \(5\) hidden layers using the LMC algorithm (note also that the autocorrelation times depend on the sampling algorithm at hand). It is worth mentioning that the blocking method can also be used to preliminarily probe the convergence properties of the simulation towards the equilibrium distribution. Typically, if the algorithm is unable to explore the configuration space consistently, for example being trapped in a sub-branch of the loss landscape, the blocking method fails, meaning that the plateau cannot be reached even for very large block sizes due to slow propagation of modes in the sampling time. We report in Fig. 16 an example of a poorly converged LMC simulation: in panel (b) it is shown that a plateau cannot be reached, indicating sampling issues.

We also assessed the convergence of Bayesian sampling using multiple parallel chains, i.e. simulating two independent Monte Carlo histories obtained with the same algorithm, but with different initial configurations and different seeds for the random number generator. In particular, we computed the Gelman–Rubin statistic \(\hat{R}\) [86], which allows to quantify the convergence of Monte Carlo simulations. Here we computed \(\hat{R}\) for the predictor of each test example and reported the averaged value. A value of the Gelman–Rubin coefficient \(\hat{R}\approx1\) signals robust convergence properties of the simulations, as shown in the representative example in Fig. 14, panel (c). On the contrary, in Fig. 16 (c), the value \(\hat{R}\approx 1.4\) signals poorly converged simulations, in agreement with what we discussed above using the blocking method. Also the blocking method benefits from multiple chain simulations: since the autocorrelation time depends only on the stochastic evolution implemented and not on the initial configuration or the random number generator, different chains should exhibit the same plateau. This is again the case in Fig. 14 (b), where the two plateau are indistinguishable up to statistical noise, while in Fig. 16 (b) the two plateau are inconsistent across the two chains. An example of the metastability effects discussed in Sec. 6.4 is reported in Fig. 15. In this case, in panel (b), it can been seen that when restricting to subsets of samples belonging to the two different states, the blocking method shows the same plateau and does not indicate that the simulations are far from equilibrium. In other words, the two slices of sampling histories appear as locally stable macrostates.

12 Details on datasets↩︎

MNIST and CIFAR10. The datasets are filtered for \(P\) training and \(P_t=1000\) test samples of the selected categorical labels, {“0”,“1”} for MNIST and {“cars”,“planes”} for CIFAR10, and the labels mapped to the numerical values \(y^\mu \in\{0,1\}\). In the case of CIFAR10, images where downsampled from \(32\times 32\) to \(28\times 28\) using bilinear interpolation with torchvision.transforms.Resize() and grayscaled, increasing slightly the difficulty of the task and resulting in the same \(N_0=784\) input dimension as for MNIST. The samples are then flattened and standardized by mean and standard deviation of the training set. To be precise, let \(\mathcal{D}_{P}=\{ x^\mu, y^\mu\}_{i=1}^{P}\) denote the training dataset, where each \(x^\mu \in \mathbb{R}^{N_0}\) is a flattened and grayscale input sample, \(y^\mu\in\mathbb{R}\), and let \(\mathcal{D}^{\text{test}}_{P_t}=\{ \tilde{x}^\mu, \tilde{y}^\mu\}_{i=1}^{P_t}\) denote the corresponding test dataset. The empirical mean and standard deviation are computed exclusively on the training set: \[m_{\mathcal{D}} = \frac{1}{P N_0} \sum_{\mu=1}^{P} \sum_{i_0=1}^{N_0} x_{i_0}^\mu, \qquad \sigma_{\mathcal{D}} = \sqrt{\frac{1}{P N_0} \sum_{\mu=1}^{P} \sum_{i_0=1}^{N_0} (x_{i_0}^\mu - m_{\mathcal{D}})^2}.\] The standardization is then applied to both training and test datasets using the same statistics: \[x^\mu \to \frac{x^\mu - m_{\mathcal{D}}}{\sigma_{\mathcal{D}}} \quad \forall\;x^\mu\in\mathcal{D}_P, \qquad \tilde{x}^\mu \to \frac{\tilde{x}^\mu - m_{\mathcal{D}}}{\sigma_{\mathcal{D}}} \quad \forall\;\tilde{x}^\mu\in\mathcal{D}_{P_t}^{\text{test}}.\] This procedure ensures that the training data has global zero mean and unit variance, while the test data is transformed consistently using the training statistics, thereby avoiding any information leakage.

Gaussian dataset. In addition to real datasets, we consider a synthetic linear regression task constructed as in classic teacher–student frameworks. Let again \(N_0\) denote the input dimensionality, which in the experiments shown in the main text is set to \(N_0 = 300\). The training inputs are independently drawn from an isotropic Gaussian distribution: \[x^\mu \sim \mathcal{N}(0, \mathbb{1}_{N_0} ), \qquad \mu = 1, \dots, P,\] and analogously for the test inputs \(\{\tilde{x}^\mu\}_{\mu=1}^{P_t}\). A fixed teacher vector \(w^\ast \in \mathbb{S}^{N_0}\) is drawn from the unit sphere by sampling and normalizing a Gaussian vector \(w^* \sim \mathcal{N}(0, \mathbb{1}_{N_0}), \;\; w^* \rightarrow w^*/\|w^*\|\). The corresponding labels are generated through the noiseless linear rule defined by the teacher vector: \[y^\mu = (x^\mu)^\top w^*, \qquad \tilde{y}^\mu = (\tilde{x}^\mu)^\top w^* \, .\] By construction, all \(x^\mu_i\) and \(y^\mu\) are standardized, and the marginal distributions of the labels inherit the Gaussian statistics of the inputs; since \(\|w^*\| = 1\), it follows that \(y^\mu \sim \mathcal{N}(0,1)\) and similarly for the test labels. The important point to note here is that there is no match between the teacher, a linear model, and the student architecture, a nonlinear MLP of depth \(L\). Therefore, even though the task is a simple linear regression, from the perspective of teacher-student models it is a non-trivial task for the student MLP to solve, which suffers from prior mismatch.

13 Numerical computation of the saddle-points↩︎

The code for theory and sampling experiments is available at Ref. [63]. Finding saddle-points of the low-dimensional effective action is possible without significant difficulties, and requires negligible compute compared to the sampling experiments.

Single-output theory. At non-zero temperature, evaluating the effective action Eq. 35 requires to compute the inverse \(y^\top \big[ \beta^{-1} \mathbb{1}+ K^{(R)}_{\mathcal{Q}}\big]^{-1} y\) each time that \(\mathcal{Q} =\prod_\ell q_\ell\) is changed, which determines the computational complexity of finding the saddle-points \(q_\ell^\ast\). Note that at zero temperature the order parameters can be pulled out as \(\mathcal{Q}^{-1} y^\top \,\big[\Theta^L(C)\big]^{-1} y\) and the contraction of the kernel inverse computed once can be reused.

We implemented the kernel functions and effective action in PyTorch, where the most efficient routine to obtain a projection of an inverse matrix \(K^{-1}y\) is torch.linalg.solve(), and also gradients of the action can be computed automatically. However, optimizing the runtime of the action evaluation was overall not necessary here. First we note that due to the absence of conjugate fields the stationary points of the action Eq. 35 are indeed minima and never saddle-points, which simplifies the optimization. We used scipy.optimize.fsolve() to perform the \(L\)-dimensional minimization over \({q_1 \dots q_L}\) in the general case, with initialization \(q_\ell =1\; \forall \ell\) at the infinite-width value. Since in the settings shown here the widths \(N_\ell\) are the same in all layers, also the action becomes symmetric and it is possible to optimize over \(q_\ell = q\; \forall \ell\), reducing to a one-dimensional minimization. It can be shown that a unique minimum exists (at finite temperature where the kernel is always invertible). The minimum can be found using a fixed scipy.optimize.brentq bracket in log space \([\log q_\mathrm{min}, \log q_\mathrm{max}]\) with \(q_\mathrm{min}=10^{-2}\), \(q_\mathrm{max}=10^{-3}\), and on the strictly monotonous transform \(\mathrm{arcsinh}\big(S(e^{\log q})\big)\), which proved to converge very fast and reliably also for the \(\mu\)P settings where \(q^\ast \gg1\). The hyperbolic arcsine and log are used here to improve the conditioning of \(S(q)\), as the minimum can be sharp. Due to \(q>0\) the \(\log q\) always exists, while \(S(q)\) can be negative and therefore \(\mathrm{arcsinh}(S)\) is used.

Non-central theory for one hidden layer (App. 9). The non-central theory implementing Eq. 57 was only needed for Fig. 9. The action Eq. 57 contains a conjugate variable \(\bar Q\), and optimization is two dimensional with the stationary points being saddles and no longer minima. We used scipy.optimize.fsolve() with initialization at the infinite-width values \(Q_0=1,\,\bar Q_0 =0\). Since the temperature is zero, the quantities defined in Eqs. 58 59 need to be computed only once.

CNN theory. For networks with convolutional layers, we optimized the effective action by performing gradient-based optimization using the Adam optimizer torch.optim.Adam. To enforce a symmetric matrix \(\mathcal{Q}\) in the optimization, it is reshaped into a vectorized form containing only the true free degrees of freedom, and the derivatives are computed with respect to this vector. Starting from an initial vectorized guess, the optimization loop continues until either the maximum number of epochs is reached or the variation in the loss function remains below a prescribed tolerance threshold for several consecutive iterations. In our simulations, we used a learning rate of \(5 \times 10^{-4}\), a tolerance of \(10^{-6}\), and a maximum number of epochs equal to \(1000\), which provided converged results.

14 Interpretation of the EWA saddle-point at zero temperature↩︎

Here we describe how the saddle-point \(q_1^\ast \dots q_L^\ast\) for single output MLPs behaves depending the relation between task and kernel structure. The interpretation is especially clear at zero temperature, due to a simplification of the minimization of the action. We start with the one-hidden layer case, where an explicit solution is available. For a lucid discussion in the deep linear network case, see also [23] Suppl.Sects. I+II.

14.1 Zero-temperature solution of the saddle-point for \(L=1\)↩︎

At zero temperature and \(L=1\), the effective action Eq. 35 is with \(q:=q_1\) \[S(q) = q - \log q + \frac{\alpha}{P} \log \det \left[\beta K^{(\text{R})}_{q}\right] + \frac{\alpha}{P} y^\top \left[K^{(\text{R})}_{q}\right]^{-1}y.\] With the definition of the scalar overlap \(M_{yy}:= \frac{1}{P}\, y^\top \Theta(C)^{-1} y\) (compare the discussion of the non-central case around Eq. 60 ), this becomes \[S(q) = q + (\alpha - 1) \log(q) + \alpha\frac{1}{q} M_{yy} +\mathrm{const.} \label{eq:app951hlzerotemp95action}\tag{68}\] First, observe that in the EWA action the only term through which the data influence the saddle point is the scalar \(M_{yy}\). This is because, while \(\log \det \Theta(C)\) is data-dependent, it is a constant offset that does not depend on \(q\). There is indeed a single minimum for all \(\alpha, M_{yy} > 0\), both positive scalars by construction, given explicitly by the solution of a quadratic equation: \[q^\ast = -\frac{\alpha - 1}{2} + \sqrt{\left(\frac{\alpha - 1}{2}\right)^2 + \alpha M_{yy}}. \label{eq:app951hlzerotemp95saddlepoint}\tag{69}\] Only the positive branch of solutions is physical due to \(q>0\), a physical constraint which is explicit in the log, but arises from the introduction of \(Q=Nq\) as a \(\mathcal{\chi}_N^2\)-distributed variable. As expected from Eq. 68 , we see that the saddle point is always growing with an increase in the task-kernel overlap \(M_{yy}\).

Note that the vicinity of \(\alpha = 1\) seems at first glance prone to produce non-monotonic behavior, also because in Eq. 68 the log term changes sign and thus pulls the solution either towards bigger or smaller \(q\), respectively. However it turns out this is not the case, instead the solution is strictly monotonically increasing with \(\alpha\) for \(M_{yy}>1\), or monotonically decreasing for \(M_{yy}<1\), as can be shown easily by calculating the derivative. The value of the solution \(q^\ast\) indeed varies monotonically with the load \(\alpha\) between the infinite-width value \(q^\ast=1\) for \(\alpha=0\) and the limiting value \(q^\ast \to M_{yy}\) for \(\alpha \to \infty\). The behavior of Eq. 69 for \(M_{yy}\in\{1/5,\,1,\,5\}\) is shown in Fig. 22 (a).

Interpretation in terms of generalization properties. Since in the zero-temperature case considered here, \(q^\ast\) cancels out in the expression for the mean predictor Eq. 48 , while the predictor variance Eq. 49 is proportional to \(q^\ast\), we obtain the following picture: Compared to the infinite-width limit at zero temperature, in the proportional regime the EWA for shallow networks of \(L=1\) predicts a change in the generalization error by modifying the predictor variance \(\mathrm{Var}(\hat{y}_{0}) \propto q^\ast\). If \(M_{yy} > 1\), then \(q^\ast > 1\) and the generalization error is increased compared to the infinite-width limit. If \(M_{yy} < 1\), then \(q^\ast < 1\) and the generalization error is decreased instead. This behavior becomes more pronounced the larger \(\alpha\) is, with a maximum factor \(M_{yy}\) multiplying the infinite-width predictor variance at \(\alpha \gg 1\) in both cases.

14.2 Behavior of the zero-temperature saddle-point for general depth \(L\)↩︎

The phenomenology of the one hidden layer action extends with very similar conclusions to the deep case, even though the resulting polynomial equation for \(\mathcal{Q}^\ast\) does not have a simple explicit solution.

For an \(L\)-layer MLP with single output dimension and rectangular aspect \(N_\ell = N \quad \forall \ell\), the minimum of the action Eq. 35 at zero temperature becomes equivalent to the minimum of \[S_V[q] = L\, q + L(\alpha - 1) \log(q) + \alpha\frac{1}{q^L} M_{yy} +\mathrm{const.} \label{eq:app95Lhl95zerotemp95action}\tag{70}\] Here we used that \(q^\ast_\ell = q^\ast \; \forall \ell\) due to the contraction principle, see Eq. 31 , and the task-kernel overlap is now \(M_{yy}=\frac{1}{P} y^\top \Theta^L(C)\,y\). Compared to Eq. 68 , the roots of the derivative are no longer determined by a quadratic function - instead we have the equation of state \[\begin{align} \left. \frac{\partial S}{\partial q} \right|_{q^\ast} \overset{!}{=} 0 \qquad \Rightarrow \qquad \begin{cases} q^\ast = 1, &\quad \mathrm{if}\; \alpha M_{yy} = 0 \\ (q^\ast)^L\, (q^\ast + \alpha - 1) = \alpha M_{yy} & \quad \mathrm{if}\; \alpha M_{yy} \neq 0 \end{cases} \label{eq:app95Lhl95zerotemp95eq-of-state} \end{align}\tag{71}\] Due to the insolvability of the quintics established by the Abel-Ruffini Theorem, for \(L\geq 5\) there is in general no explicit solution to this equation. Nonetheless the implicit solution \[q^\ast = \left( \frac{\alpha M_{yy}}{q^\ast + \alpha - 1} \right)^{\frac{1}{L}} \label{eq:app95Lhl95zerotemp95saddlepoint95implicit}\tag{72}\] allows us to show that the qualitative phenomenology of the \(L=1\) case is preserved. First, this equation still always has a unique solution: The positive half-space of the RHS (where \(q^\ast > 1 - \alpha\)) is a function decreasing monotonically from \(+\infty\) to \(0\), such that for any \(\alpha\) there always is exactly one intersection with the LHS, and the negative half-space of the RHS can never intersect with the LHS for physical values \(q^\ast > 0\).

Second, we find that \(\bar{Q}_\ast = 1\) for \(\alpha = 0\), \(\bar{Q}_\ast = \sqrt[L+1]{M_{yy}}\) for \(\alpha = 1\), and \(\bar{Q}_\ast = \sqrt[L]{M_{yy}}\) for \(\alpha \to \infty\); and monotonicity can again be confirmed by taking a derivative of the RHS. Taking into account that for the deep case the factor multiplying the kernel, and correspondingly also the predictor variance, is \(\mathcal{Q}^\ast = (q^\ast)^L\) we find qualitatively similar behavior as for \(L=1\) also in the deep case, shown in Fig. 22 (c).

Third, we ask how depth changes the size of \(\mathcal{Q}^\ast\) at intermediate values of \(\alpha\). Note that Eq. 72 can be solved for \(\alpha\) as \[\alpha(q^\ast,M_{yy},L) = \frac{q^\ast - 1}{(q^\ast)^{-L} M_{yy} \; - 1};\] plugging in the value \((q^\ast_h)^L = (M_{yy}+1)/2\) halfway between the two limiting values \(\{1, M_{yy}\}\), this gives to first order in \(1/L\) \[\alpha(q^\ast_h,M_{yy},L) \propto \exp\left[\frac{1}{L} \log\bigg(\frac{M_{yy}+1}{2}\bigg)\right] -1 = \frac{1}{L} \log\bigg(\frac{M_{yy}+1}{2}\bigg) + O(L^{-2}).\] Interpretation in the \(L\)-layer case. As for the shallow network, also in the deep case for \(M_{yy}>1\) the generalization error is increased, and for \(M_{yy}<1\) decreased with respect to the infinite-width limit. This is because the predictor variance \(\mathrm{Var}(\hat{y}_{0}) \propto \mathcal{Q}^\ast\), with a maximum factor \(\mathcal{Q}^\ast\to M_{yy}\) for \(\alpha\gg1\). The difference to the shallow case is that depth accelerates the transition to larger factors \(\mathcal{Q}^\ast\) and therefore the effect of finite-width: The value \(\alpha_h\) where \(\mathcal{Q}^\ast\) is halfway between the limiting values \(\{1, M_{yy}\}\) scales with depth as \(\alpha_h \propto 1/L\). This again implies that the product \(L\alpha\) is the effective load parameter controlling the size of finite-width effects in deep networks.

References↩︎