Asymptotics of representation learning in finite Bayesian neural networks
Year: 2022

Abstract:
Recent works have suggested that finite Bayesian neural networks may sometimes outperform their infinite cousins because finite networks can flexibly adapt their internal representations. However, our theoretical understanding of how the learned hidden layer representations of finite networks differ from the fixed representations of infinite networks remains incomplete. Perturbative finite-width corrections to the network prior and posterior have been studied, but the asymptotics of learned features have not been fully characterized. Here, we argue that the leading finitewidth corrections to the average feature kernels for any Bayesian network with linear readout and Gaussian likelihood have a largely universal form. We illustrate this explicitly for three tractable network architectures: deep linear fully-connected and convolutional networks, and networks with a single nonlinear hidden layer. Our results begin to elucidate how task-relevant learning signals shape the hidden layer representations of wide Bayesian neural networks.

heading: 1 Introduction
text: The expressive power of deep neural networks critically depends on their ability to learn to represent the features of data [1–24]. However, the structure of their hidden layer representations is only theoretically well-understood in certain infinite-width limits, in which these representations cannot flexibly adapt to learn data-dependent features [3–11, 24]. In the Bayesian setting, these representations are described by fixed, deterministic kernels [3–11]. As a result of this inflexibility, recent works have suggested that finite Bayesian neural networks (henceforth BNNs) may generalize better than their infinite counterparts because of their ability to learn representations [10].
Theoretical exploration of how finite and infinite BNNs differ has largely focused on the properties of the prior and posterior distributions over network outputs [12–17]. In particular, several works have studied the leading perturbative finite-width corrections to these distributions [12–16]. Yet, the corresponding asymptotic corrections to the feature kernels, which measure how representations evolve from layer to layer, have only been studied in a few special cases [16]. Therefore, the structure of these corrections, as well as their dependence on network architecture, remain poorly understood. In this paper, we make the following contributions towards the goal of a complete understanding of feature learning at asymptotically large but finite widths:
• We argue that the leading finite-width corrections to the posterior statistics of the hidden layer kernels of any BNN with a linear readout layer and Gaussian likelihood have a largely
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
prescribed form (Conjecture 1). In particular, we argue that the posterior cumulants of the kernels have well-defined asymptotic series in terms of their prior cumulants, with coefficients that have fixed dependence on the target outputs.
• We explicitly compute the leading finite-width corrections for deep linear fully-connected networks (§4.1), deep linear convolutional networks (§4.2), and networks with a single nonlinear hidden layer (§4.3). We show that our theory yields quantitatively accurate predictions for the result of numerical experiment for tractable linear network architectures, and qualitatively accurate predictions for deep nonlinear networks, where quantitative analytical predictions are intractable.
Our results begin to elucidate the structure of learned representations in wide BNNs. The assumptions of our general argument are satisfied in many regression settings, hence our qualitative conclusions should be broadly applicable.

heading: 2 Preliminaries
text: We begin by defining our notation, setup, and assumptions. We will index training and test examples by Greek subscripts µ, ν, . . ., and layer dimensions (that is, neurons) by Latin subscripts j, l, . . .. Layers will be indexed by the script Latin letter `. Matrix- or vector-valued quantities corresponding to a given layer will be indexed with a parenthesized superscript, while scalar quantities that depend only on the layer will be indexed with a subscript. Depending on context, ‖ · ‖ will denote the `2 norm on vectors or the Frobenius norm on matrices. We denote the standard Euclidean inner product of two vectors a,b ∈ Rn by a · b.

heading: 2.1 Bayesian neural networks with linear readout
text: Throughout this paper, we consider deep Bayesian neural networks with fully connected linear readout. Such a network f : Rn0 → Rnd with d layers can be written as
f(x;W d,W) = 1√ nd−1 W (d)ψ(x;W), (1)
where the feature map ψ(·;W) : Rn0 → Rnd−1 includes all d − 1 hidden layers, collectively parameterized byW . Here, ψ can be some combination of fully-connected feedforward networks, convolutional networks, recurrent networks, et cetera; we assume only that it has a well-defined infinite-width limit in the sense of §2.2. We let the widths of the hidden layers be n1, n2, . . . , nd−1; we define the width of a convolutional layer to be its channel count [7]. We assume isotropic Gaussian priors over the trainable parameters [1–23], with W (d)ij ∼i.i.d N (0, σ2d) in particular.
In our analysis, we fix an arbitrary training dataset D = {(xµ,yµ)}pµ=1 of p examples. We define the input and output Gram matrices of this dataset as [Gxx]µν ≡ n−10 xµ · xν and [Gyy]µν ≡ n −1 d yµ · yν , respectively. For analytical tractability, we consider a Gaussian likelihood p(D |Θ) ∝ exp(−βE) for
E(Θ;D) = 1 2 p∑ µ=1 ‖f(xµ; Θ)− yµ‖2, (2)
where β ≥ 0 is an inverse temperature parameter that sets the variance of the likelihood and Θ = {W (d),W} [23]. We then introduce the Bayes posterior over parameters given these data:
p(Θ | D) = p(D |Θ)p(Θ) p(D) ; (3)
we denote averages with respect to this distribution by 〈·〉. By tuning β, one can then adjust whether the posterior is dominated by the prior (β 1) or the likelihood (β 1). We will mostly focus on the case in which the input dimension is large and the training dataset can be linearly interpolated; the low-temperature limit β →∞ then enforces the interpolation constraint.

heading: 2.2 The Gaussian process limit
text: We consider the limit of large hidden layer widths n1, n2, . . . , nd−1 →∞ with n0, nd, p, and d fixed. More precisely, we consider a limit in which n` = α`n for ` = 1, . . . , d− 1, where α` ∈ (0,∞) and
n→∞, as studied by [3–15, 17–19, 24] and others. Importantly, we note that size of n0 relative to n is unimportant for our results, whereas nd/n and d/n must be small [10, 12, 17].
In this limit, for ψ built out of compositions of most standard neural network architectures, the prior over function values f tends to a Gaussian process (GP) [3–8]. Moreover, with our choice of a Gaussian likelihood, the posterior over function values also tends weakly to the posterior induced by the limiting GP prior [25]. The kernel of the limiting GP prior is given by the deterministic limit K (d−1) ∞ of the inner product kernel of the postactivations of the final hidden layer,
K(d−1)(x,x′) ≡ n−1d−1ψ(x,W) ·ψ(x ′,W), (4)
multiplied by the prior variance σ2d [3–8]. For a broad range of network architectures, K (d−1) ∞ can be computed recursively [5–8]. For brevity, we define the kernel matrix evaluated on the training data: [K(d−1)]µν ≡ K(d−1)(xµ,xν).

heading: 3 Elementary perturbation theory for finite Bayesian neural networks
text: We first present our main result, which shows that the form of the leading perturbative correction to the average hidden layer kernels of a BNN is tightly constrained by the assumptions that the readout is linear, that the cost is quadratic, and that the GP limit is well-defined.

heading: 3.1 Finite-width corrections to the posterior cumulants of hidden layer observables
text: Our main result is as follows:
Conjecture 1 Consider a BNN of the form (1), with posterior (3). Assume that this network admits a well-defined GP limit as discussed in §2.2. Let O be a hidden layer observable, that is, a function of the hidden layer activations that is not a function of the readout weights Wd. Assume that O tends in probability to a finite, deterministic limit O∞ under the posterior in the GP limit.
Then, the posterior cumulants of this observable admit well-behaved asymptotic series at large widths in terms of its joint prior cumulants with the postactiviation kernel K(d−1). In particular, the asymptotic expansion of the posterior mean 〈O〉 has leading terms
〈O〉 = EWO + 1
2 nd p∑ ρ,λ=1 [σ−2d Γ −1GyyΓ −1 − Γ−1]ρλ covW(O,K(d−1)ρλ ) + . . . , (5)
where Γ ≡ K(d−1)∞ + β−1σ−2d Ip. Here, the cumulants of the kernels are computed with respect to the prior, and are themselves given by asymptotic series at large widths. The ellipsis denotes terms that are of subleading order in the inverse hidden layer widths.
In Appendix B, we derive this result perturbatively by expanding the posterior cumulant generating function of O in powers of the deviations of O and K(d−1) from their deterministic infinite-width values. There, we also give an asymptotic formula for the posterior covariance of two observables. However, the resulting perturbation series may not rigorously be an asymptotic series, and this method does not yield quantitative bounds for the width-dependence of the terms. We therefore frame it as a conjecture. We note that similar methods can be applied to compute asymptotic corrections to the posterior predictive statistics; we comment on this possibility in Appendix G.
Though this conjecture applies to a broad class of hidden layer observables, the observables of greatest interest are the preactivation or postactivation kernels of the hidden layers within the feature map ψ. We will focus on the postactivation kernels K(`), which measure how the similarities between inputs evolve as they are propagated through the network [5–10].
Conjecture 1 posits that there are two possible types of leading finite-width corrections to the average kernels. The first class of corrections are deviations of EWK(`) from K(`)∞ . These terms reflect corrections to the prior, and do not reflect non-trivial representation learning as they are independent of the outputs. For fully-connected networks, also known as multilayer perceptrons (MLPs), work by Yaida [12] and by Gur-Ari and colleagues [18, 19] shows that EWK(`) = K(`)∞ +O(n−1). The
second type of correction is the output-dependent term that depends on covW(K (`) µν ,K (d−1) ρλ ). For deep linear MLPs or MLPs with a single hidden layer, EWK(`) is exactly equal to K(`)∞ at any width (see Appendix C) [3, 12, 18], and only the covariance term contributes. More broadly, these prior works show that covW(K (`) µν ,K (d−1) ρλ ) = O(n−1) for MLPs, and that higher cumulants are of O(n−2) [12, 18, 19]. Some of these results have recently been extended to convolutional networks by Andreassen and Dyer [26]. Thus, the finite-width correction to the prior mean should not dominate the feature-learning covariance term, and the terms hidden in the ellipsis should indeed be suppressed.
The leading output-dependent correction has several interesting features. First, it includes a factor of nd, reflecting the fact that inference in wide Bayesian networks with many outputs is qualitatively different from that in networks with few outputs relative to their hidden layer width [10]. If nd/n does not tend to zero with increasing n, the infinite-width behavior is not described by a standard GP [8, 10]. Moreover, we note that the matrix Γ is invertible at any finite temperature, even when K(d−1)∞ is singular. Therefore, provided that one can extend the GP kernel by continuity to non-invertible Gxx, Conjecture 1 can be applied in the data-dense regime n0 < p as well as the data-sparse regime n0 > p. Furthermore, we observe that the correction depends on the outputs only through their Gram matrix Gyy . This result is intuitively sensible, since with our choice of likelihood and prior the function-space posterior is invariant under simultaneous rotation of the output activations and targets. Finally, Gyy is transformed by factors of the matrix Γ−1, hence the correction depends on certain interactions between the output similarities and the GP kernel K(d−1)∞ .

heading: 3.2 High- and low-temperature limits of the leading correction
text: To gain some intuition for the properties of the leading finite-width corrections, we consider their high- and low-temperature limits. These limits correspond to tuning the posterior (3) to be dominated by the prior or the likelihood, respectively. At high temperatures (β 1), expanding Γ−1 as a Neumann series (see Appendix A and [27]) yields
σ−2d Γ −1GyyΓ −1 − Γ−1 = −βσ2dIp + (βσ2d)2(σ−2d Gyy +K (d−1) ∞ ) +O[(βσ2d)3]. (6)
Thus, at high temperatures, the outputs only influence the average kernels of Conjecture 1 to subleading order in both width and β, which reflects the fact that the likelihood is discounted relative to the prior in this regime. Moreover, the leading output-dependent contribution averages together Gyy and K (d−1) ∞ , hence, intuitively, there is no way to ‘cancel’ the GP contributions to the average kernels. We note that, at infinite temperature (β = 0), the posterior reduces to the prior, and all finite-width corrections to the average kernels arise from the discrepancy between EWK(`) and K(`)∞ .
At low temperatures (β 1), the behavior of Γ−1 differs depending on whether or not K(d−1)∞ is of full rank. Assuming for simplicity that it is invertible, we have
σ−2d Γ −1GyyΓ −1 − Γ−1 = [K(d−1)∞ ]−1(σ−2d Gyy −K (d−1) ∞ )[K (d−1) ∞ ] −1 +O[(βσ2d)−1]; (7)
in the non-invertible case there are additional contributions involving projectors onto the null space of K (d−1) ∞ . Therefore, the leading-order low temperature correction depends on the difference between the target and GP kernels, while the leading non-trivial high temperature correction depends on their sum.

heading: 4 Learned representations in tractable network architectures
text: Having derived the general form of the leading perturbative finite-width correction to the average feature kernels, we now consider several example network architectures. For these tractable examples, we provide explicit formulas for the feature-learning corrections to the hidden layer kernels, and test the accuracy of our theory with numerical experiments.

heading: 4.1 Deep linear fully-connected networks
text: We first consider deep linear fully-connected networks with no bias terms. Concretely, we consider a network with activations h(`) ∈ Rn` recursively defined via h(`) = n−1/2`−1 W (`)h(`−1) with base case
h(0) = x, where the prior distribution of weights is [W (`)]ij ∼i.i.d. N (0, σ2` ). For such a network, the hidden layer kernels [K(`)]µν ≡ n−1` h (`) µ · h(`)ν have deterministic limits K(`)∞ = m2`Gxx, where m2` ≡ σ2`σ2`−1 · · ·σ21 is the product of prior variances up to layer `. Higher prior cumulants of the kernels are easy to compute with the aid of Isserlis’ theorem for Gaussian moments (see Appendix C) [28, 29], yielding
〈K(`)〉 m2` = Gxx + ( ∑̀ `′=1 nd n`′ ) GxxΓ −1 (m−2d Gyy − Γ)Γ−1Gxx +O(n−2), (8) where Γ ≡ Gxx + Ip/(βm2d) and ` = 1, . . . , d− 1. In Appendix D, we show that this result can be derived directly through an ab initio perturbative calculation of the cumulant generating function of the kernels, without relying on our heuristic argument for the general version of Conjecture 1. Moreover, in Appendix E, we show that the form of the correction remains the same even if one allows arbitrary forward skip connections, though the dependence on width and depth is given by a more complex recurrence relation.
Thus, the leading corrections to the normalized average kernels 〈K(`)〉/m2` are identical across all hidden layers up to a scalar factor that encodes the width-dependence of the correction. This sum-of-inverse-widths dependence was previously noted by Yaida [12] in his study of the corrections to the prior of a deep linear network. For a network with hidden layers of equal width n, we have the simple linear dependence ∑` `′=1(nd/n`′) = nd`/n. If one instead includes a narrow bottleneck in an otherwise wide network, this dependence predicts that the kernels before the bottleneck should be close to their GP values, while those after the bottleneck should deviate strongly.
This result simplifies further at low temperatures, where, by the result of §3.2, we have
〈K(`)〉 m2` = Gxx + ( ∑̀ `′=1 nd n`′ )( m−2d Gyy −Gxx ) +O(n−2, β−1) (9)
in the regime in which Gxx is invertible. We thus obtain the simple qualitative picture that the low-temperature average kernels linearly interpolate between the input and output Gram matrices. In Appendix F, we show that this limiting result can be recovered from the recurrence relation derived through other methods by Aitchison [10], who did not use it to compute finite-width corrections. We note that the low-temperature limit is peculiar in that the mean predictor reduces to the least-norm pseudoinverse solution to the underlying underdetermined linear system XW = Y ; we comment on this property in Appendix G.
We can gain some additional understanding of the structure of the correction by using the eigendecomposition of Gxx. As Gxx is by definition a real positive semidefinite matrix, it admits a unitary eigendecomposition Gxx = UΛU† with non-negative eigenvalues Λµµ. In this basis, the average kernel is
1
m2` U†〈K(`)〉U = Λ + ( ∑̀ `′=1 nd n`′ )( m−2d Λ̃U †GyyU Λ̃− Λ̃Λ ) +O(n−2), (10)
where we have defined the diagonal matrix Λ̃ ≡ βm2dΛ(Ip+βm2dΛ)−1. As βm2dΛ ≥ 0, the diagonal elements of Λ̃ are bounded as 0 ≤ Λ̃µµ ≤ 1. Thus, the factors of Γ−1Gxx by whichGyy is conjugated have the effect of suppressing directions in the projection of Gyy onto the eigenspace of Gxx with small eigenvalues. We can see that this effect will be enhanced at high temperatures (β 1) and small scalings (m2d 1), and suppressed at low temperatures and large scalings. For this linear network, similarities are not enhanced, only suppressed. Moreover, if Gxx is diagonal, then a given element of the average kernel will depend only on the corresponding element of Gyy .
We now seek to numerically probe how accurately these asymptotic corrections predict learned representations in deep fully-connected linear BNNs. Using Langevin sampling [30, 31], we trained deep linear networks of varying widths, and compared the difference between the empirical and GP kernels with theory predictions. We provide a detailed discussion of our numerical methods in Appendix I. In Figure 1, we present an experiment with a 2-layer linear neural network trained on the MNIST dataset of handwritten digit images [32] using the Neural Tangents library [33]. We find an excellent agreement with our theory, confirming the inverse scaling with width and linear scaling with depth for the deviations from GP kernel.

heading: 4.2 Deep linear convolutional networks
text: To demonstrate the applicability of Conjecture 1 to non-fully-connected BNNs, we consider deep convolutional linear networks with no bias terms. Here, the appropriate notion of width is the number of channels in each hidden layer [7]. Following the setup of Novak et al. [7] and Xiao et al. [34], we consider a network consisting of d − 1 linear convolutional layers followed by a fully-connected linear readout layer. For simplicity, we restrict our attention to convolutions with periodic boundary conditions, and do not include internal pooling layers (see Appendix C for more details). Concretely, we consider a network with hidden layer activations h(`)i,a, where i indexes the n` channels of the layer and a is a spatial multi-index. The hidden layer activations are then defined through the recurrence
h (`) i,a(x) = 1 √ n`−1 n`−1∑ j=1 ∑ b w (`) ij,bh (`−1) j,a+b(x) (11)
with base case h(0)i,a (x) = xi,a, where i indexes the input channels (e.g., image color channels). The feature map is then formed by flattening the output of the last hidden layer into an nd−1s-dimensional vector, where s is the total dimensionality of the inputs (see Appendix C for details). We fix the prior distribution of the filter elements to be w(`)ij,a ∼i.i.d. N (0, σ 2 ` va), where va > 0 is a weighting factor
that sets the fraction of receptive field variance at location a (and is thus subject to the constraint∑ a va = 1). For inputs [xµ]i,a and [xν ]i,a, we introduce the four-index hidden layer kernels
K (`) µν,ab ≡
1
n` n∑̀ i=1 h (`) i,a(xµ)h (`) i,b(xν). (12)
With the given readout strategy, the two-index feature map kernel appearing in Conjecture 1 is related to the four-index kernel of the last hidden layer by K(d−1)µν = 1s ∑ aK (d−1) µν,aa . We discuss other readout strategies in Appendix C, but use this vectorization strategy in our numerical experiments.
As shown by Xiao et al. [34], the infinite-width four-index kernel obeys the recurrence
[K(`)∞ ]µν,ab = σ 2 ` ∑ c vc[K (`−1) ∞ ]µν,(a+c)(b+c) (13)
with base case [K0∞]µν,ab = [Gxx]µν,ab ≡ 1n0 ∑n0 i=1[xµ]i,a[xν ]i,b. This gives convolutional linear networks a sense of spatial hierarchy that is not present in the fully-connected case: even at infinite width, the kernels include iterative spatial averaging.
In Appendix C, we derive the kernel covariances appearing in Conjecture 1. As in the fully-connected case, this computation is easy to perform with the aid of Isserlis’ theorem. The general result is
somewhat complicated, but things simplify under the assumption that readout is performed using vectorization. Then, one finds that
〈K(`)µν,ab〉 = [K (`) ∞ ]µν,ab+ ( d−1∏ `′=` σ2` )(∑̀ `′=1 nd n`′ ) 1 s s∑ c=1 p∑ ρ,λ=1 [K(`)∞ ]µρ,acΦρλ[K (`) ∞ ]λν,cb+O(n−2), (14) where we have defined Φρλ ≡ [σ−2d Γ−1GyyΓ−1 − Γ−1]ρλ for brevity. Thus, the correction to the convolutional kernel is quite similar to that obtained in the fully-connected case. To this order, the difference between these network architectures manifests itself largely through the difference in the infinite-width kernels. In Appendix C, we show that a similar simplification holds if readout is performed using global average pooling over space.
As we did for fully-connected networks, we test whether our theory accurately predicts the results of numerical experiment, using the MNIST digit images illustrated in 2(a-d). We consider a network with one-dimensional (Figure 2e and f) and two-dimensional (Figure 3) convolutional hidden layers, trained to classify 50 MNIST images (see Appendix I for details of our numerical methods). As shown in Figure 2(e, f) (Figure 3(a,b) for 2D convolutions), we again obtain good quantitative agreement between the predictions of our asymptotic theory and the results of numerical experiment. In Figure 3c, we directly visualize the learned feature kernels for 2D convolutional layers, illustrating the good agreement between theory and experiment. Therefore, our asymptotic theory can be applied to accurately predict learned representations in deep convolutional linear networks.

heading: 4.3 Networks with a single nonlinear hidden layer
text: Finally, we would like to gain some understanding of how including nonlinearity affects the structure of learned representations. However, for a nonlinear MLP, it is usually not possible to analytically compute covW(K (`) µν ,K (d−1) ρλ ) to the required order [9, 12, 18, 19]. Here, we consider the case of a network with a single nonlinear layer and no bias terms, in which we can both summarize the key obstacles to studying deep nonlinear networks and gain some intuitions about how they might differ from linear BNNs. Concretely, we consider a network with feature map ψ(x;W (1)) =
φ(n −1/2 0 W
(1)x) for an elementwise activation function φ, where the weight matrix W (1) has prior distribution [W (1)]ij ∼i.i.d. N (0, σ21). The only hidden layer kernel of this network is the feature map postactivation kernel Kµν defined in (4), where we drop the layer index for brevity. As detailed in Appendix H, for such a network we have the exact expressions
[K∞]µν = EWKµν = E[φ(hµ)φ(hν)], (15) n1 covW(Kµν ,Kρλ) = E[φ(hµ)φ(hν)φ(hρ)φ(hλ)]− [K∞]µν [K∞]ρλ, (16)
where expectations are taken over the p-dimensional Gaussian random vector hµ, which has mean zero and covariance cov(hµ, hν) = σ21 [Gxx]µν . Unlike for deeper nonlinear networks, here there are no finite-width corrections to the prior expectations [3, 12, 18].
Though these expressions are easy to define, it is not possible to evaluate the four-point expectation in closed form for general Gram matrices Gxx and activation functions φ, including ReLU and erf. This obstacle has been noted in previous studies [9, 12, 15], and makes it challenging to extend approaches similar to those used here to deeper nonlinear networks. For polynomial activation functions, the required expectations can be evaluated using Isserlis’ theorem (see Appendix A). However, even for a quadratic activation function φ(x) = x2, the resulting formula for the kernel will involve many elementwise matrix products, and cannot be simplified into an intuitively comprehensible form.
If the input Gram matrix Gxx is diagonal, the four-point expectation becomes tractable because the required expectations factor across sample indices. In this simple case, there is an interesting distinction between the behavior of activation functions that yield Eφ(h) = 0 and those that yield Eφ(h) 6= 0. As detailed in §4.1, if Eφ(h) = 0, K∞ is diagonal, and a given element of the leading finite-width correction to 〈K〉 depends only on the corresponding element of Gyy. However, if Eφ(h) 6= 0, then K∞ includes a rank-1 component, and each element of the correction depends on all elements of Gyy . This means that the case in which Gxx is diagonal is qualitatively distinct from the case in which there is only a single training input for such activation functions.

heading: 5 Learned representations in deep nonlinear networks
text: In the preceding section, we noted that analytical study of learned representations in deep nonlinear BNNs is generally quite challenging. Here, we use numerical experiments to explore whether any
of the intuitions gained in the linear setting carry over to nonlinear networks. Concretely, we study how narrow bottlenecks affect representation learning in a more realistic nonlinear network. We train a network with three hidden layers and ReLU activations on a subset of the MNIST dataset [32]. Despite its analytical simplicity, ReLU is among the activation functions for which the covariance term in Conjecture 1 cannot be evaluated in closed form (see §4.3). However, it is straightforward to simulate numerically. Consistent with the predictions of our theory for linear networks, we find that introducing a narrow bottleneck leads to more representation learning in subsequent hidden layers, even if those layers are quite wide (Figure 4). Quantitatively, if one increases the width of the hidden layers between which the fixed-width bottleneck is sandwiched, the deviation of the first layer’s kernel from its GP value decays roughly as 1/n with increasing width, while the deviations for the bottleneck and subsequent layers remain roughly constant. In contrast, the kernel deviations throughout a network with equal-width hidden layers decay roughly as 1/n (Figure 4). These observations are qualitatively consistent with the width-dependence of the linear network kernel (8), as well as with previous studies of networks with infinitely-wide layers separated by a finite bottleneck [35]. Keeping in mind the obstacles noted in §4.3, precise characterization of nonlinear networks will be an interesting objective for future work.

heading: 6 Related work
text: Our work is closely related to several recent analytical studies of finite-width BNNs. First, Aitchison [10] argued that the flexibility afforded by finite-width BNNs can be advantageous. He derived a
recurrence relation for the learned feature kernels in deep linear networks, which he solved in the limits of infinite width and few outputs, narrow width and many outputs, and infinite width and many outputs. As discussed in §4.1 and in Appendix F, our results on deep linear networks extend those of his work. Furthermore, our numerical results support his suggestion that networks with narrow bottlenecks may learn interesting features.
Moreover, our analytical approach and the asymptotic regime we consider mirror recent perturbative studies of finite-width BNNs. As noted in §3 and Appendix B, we make use of the results of Yaida [12], who derived recurrence relations for the perturbative corrections to the cumulants of the finitewidth prior for an MLP. However, Yaida did not attempt to study the statistics of learned features; the goal of his work was to establish a general framework for the study of finite-width corrections. Bounds on the prior cumulants of a broader class of observables have been studied by Gur-Ari and colleagues [18, 19, 26]; these results could allow for the identification of observables to which Conjecture 1 should apply. Finally, perturbative corrections to the network prior and posterior have been studied by Halverson et al. [13] and Naveh et al. [15], respectively. Our work builds upon these studies by perturbatively characterizing the internal representations that are learned upon inference.
Following the appearance of our work in preprint form, Roberts et al. [36] announced an alternative derivation of the zero-temperature limit of Conjecture 1 for MLPs; we have adopted their terminology of hidden layer observables. As in Yaida [12]’s earlier work, they rely on sequential perturbative approximation of the prior over preactivations as the hidden layers are marginalized out in order from the first to the last. While our elementary perturbative argument for Conjecture 1 does not require assuming a particular network architecture for the hidden layers, it takes as input information regarding the prior cumulants that would have to be approximated using such methods. Moreover, the approach of layer-by-layer approximation to the prior could enable a fully rigorous version of Conjecture 1 to be proved on an architecture-by-architecture basis [37].
Our work, like most studies of wide BNNs [3–15, 17–19, 24], focuses on the regime in which the sample size p is held fixed while the hidden layer width scale n tends to infinity, i.e., p n. One can instead consider regimes in which p is not negligible relative to n, in which the posterior would be expected to concentrate. The behavior of deep linear BNNs in this regime was recently studied by Li and Sompolinsky [16], who computed asymptotic approximations for the predictor statistics and hidden layer kernels. In Appendix F, we show that our result (9) for the zero-temperature kernel can be recovered as the p/n ↓ 0 limit of their result. As the dataset size p appears only implicitly in our approach, we leave the incorporation of large-p corrections as an interesting objective for future work. We note, however, that alternative methods developed to study the large-p regime [16, 38] cannot overcome the obstacles to analytical study of deep nonlinear networks encountered here.

heading: 7 Conclusions
text: In this paper, we have shown that the leading perturbative feature learning corrections to the infinitewidth kernels of wide BNNs with linear readout and least-squares cost should be of a tightly constrained form. We demonstrate analytically and with numerical experiments that these results hold for certain tractable network architectures, and conjecture that they should extend to more general network architectures that admit a well-defined GP limit.
Limitations. We emphasize that our perturbative argument for Conjecture 1 is not rigorous, and that we have not obtained quantitative bounds on the remainder for general network architectures. It is possible that there are non-perturbative contributions to the posterior statistics that are not captured by Conjecture 1; non-perturbative investigation of feature learning in finite BNNs will be an interesting objective for future work [17, 39]. More broadly, we leave rigorous proofs of the applicability of our results to more general architectures and of the smallness of the remainder as objective for future work. As mentioned above, one could attempt such a proof on an architecture-by-architecture basis [12, 36, 37]. Alternatively, one could attempt to treat all sufficiently sensible architectures uniformly [8, 9]. Furthermore, we have considered only one possible asymptotic regime: that in which the width is taken to infinity with a finite training dataset and small output dimensionality. As discussed above in reference to the work of Aitchison [10] and Li and Sompolinsky [16], investigation of alternative limits in which output dimension, dataset size, depth, and hidden layer width are all taken to infinity with fixed ratios may be an interesting subject for future work.
Acknowledgments and Disclosure of Funding
We thank B. Bordelon for helpful comments on our manuscript. JAZ-V acknowledges partial support from the NSF-Simons Center for Mathematical and Statistical Analysis of Biology at Harvard and the Harvard Quantitative Biology Initiative. This work was further supported by the Harvard Data Science Initiative Competitive Research Fund, the Harvard Dean’s Competitive Fund for Promising Scholarship, and a Google Faculty Research Award. The authors declare no conflict of interest.

Review summaries:
This paper shows analytically and numerically that the leading finite-width corrections to the average hidden layer kernels of any fully-connected feedforward networks with linear readout and least-squares loss have a largely prescribed form :  these assumptions fix the dependency of the correction on the target outputs, the size of the correction scale with the ratio of the number of outputs to the typical hidden layer width. Authors explicitly compute the leading finite-width corrections for deep linear networks 	and networks with a single nonlinear hidden layer . Theyshow how certain directions in the output kernel matrix can be enhanced or suppressed in the hidden layer representations depending on the structure of the input similarities   This is equivalent to considering the equilibrium distribution of Langevin sampling of the parameters at inverse temperature β. Moreover, by tuning β, one can then adjust whether the posterior is dominated by the prior (β ≪ 1) or the likelihood (β ≫ 1).    Authors conjecture, non-rigorously, that these results should extend to Bayesian neural network with different architectures  

Layer outputs of neural networks with Gaussian priors are proved to be Gaussian processes at infinite width. The paper is focused on obtaining a characterization of layer outputs at finite width. Based on the NNGP kernel resulting in the wide regime, the main result is the introduction of perturbative corrections to the average of this kernel. These corrections appear in finite-width neural networks and don’t exist in the limit.   **Contributions**:  * Updated description of Yaida [11] on the difference between finite and infinite wide neural networks  * Feature learning interpretation   **Limitations**: * Not rigorous  * Works on very simple setting but no additional experiments  * Not evident how to use the results

This paper studies the problem of leading-order perturbations to average feature kernels of Bayesian finite-but-wide neural networks (NNs). The paper is mainly of theoretical contribution and the main result establishes a conjecture for the leading order corrections for general NN architectures under linear readout and squared loss. The authors then delve into limiting situations of tempered posteriors before analysing two tractable cases: deep linear NNs and single-hidden layer non-linear NNs. Experimental evidence on toy settings support their claims.