Asymptotics of representation learning in finite Bayesian neural networks
Authors: Jacob A Zavatone-Veth
Abdulkadir Canatar
Ben Ruben
Cengiz Pehlevan
Venue: NeurIPS 2021 Poster (2021)
Keywords: Representation learning
Bayesian neural networks
finite width corrections

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 finite-width 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. 

Asymptotics of representation learning in finite Bayesian neural networks
Jacob A. Zavatone-Vethh Abdulkadir Canatar BeBjamin Ruben',
Cengiz Pehlevan
Department of Physics, Center for Brain Science, 'Biophysies Graduate Program, *John A. Paulson School of Engineering and Applied Sciences Harvard University
Cambridge, MA 02138
{j
(jzavatoneveth, canatara, benruben) hww.arrraa....... edu opehlevan@seas harvard. edu
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-wine 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 finite- width corrections to average feature kernels for any Bayesian network with linear readout and Gaussian likelihood have largely universal form. We illustrate this explicitly for three tractable network architectures: deep linear fully-connected and convolutional networks, and networks with single nonlinear hidden layer. Our results begin to elucidate how task-relevant learning signals shape the hidden layer representations of wide Bayesian neural networks.
Introduction
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 their hidden layer representations only theoretically well-understood in certain infinite-width limits, in which these representations cannot flexibly adapt learn data-dependent features [3-11, 24]. In the Bayesian setting. these represen- tations are described by fixed, deterministic kernels [3-11]. As 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 these distributions [12-16]. Yet, the corresponding asymptotic corrections to the feature kernels, which measure how representations evolve from layer layer, have only been studied in 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 complete understanding of feature learning at asymptotically large but finite widths:
We argue that the leading Enite-widtt corrections to the posterior statistics of the hidden layer kemelss of any BNN with linear readout layer and Gaussian likelihood have largely
35th Conference on Neural Information Processing Systems (NcurIPS 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 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.
2
Preliminaries
We begin by defining our notation, setup, and assumptions. We will index training and test examples by Greek subscripts μt, u, and layer dimensions (that is, neurons) by Latin subscripts j, Layers will be indexed by the script Latin letter f. Matrix- or vector-valued quantities corresponding given layer will be indexed with parenthesized superscript, while scalar quantities that depend only on the layer will be indexed with a subscript. Depending on context, || || will denote the & norm vectors or the Frobenius norm matrices. We denote the standard Euclidean inner product of two vectors a, b E R by a b.
2.1 Bayesian neural networks with linear readout
Throughout this paper, we consider deep Bayesian neural networks with fully connected linear readout. Such a network f: R + gnd with d layers can be written as
f(x: Wd. W)
1
W (1) (x;;;; W),
(I)
where the feature map (b; W) R + R 7 includes all d hidden layers, collectively parameterized by W. Here, &b can be some combination of fully-connected feedforward networks, convolutional networks, recurrent networks, et cetera; we assume only that it has well-defined infinite-width limit in the sense of $2.2. We let the widths of the hidden layers be n1, ,1,722 .......... fdd-1; we define the width of convolutional layer to its channel count [7]. We assume isotropic Gaussian priors over the trainable parameters [1-23], with wr(d)
~isd WO0, o.)) in particular.
In our analysis, we fix an arbitrary training dataset D = {(x,, yw) of p examples. We define the input and output Gram matrices of this dataset as (Gaz)ye no X X, and [Gyy]e = n yy Yu. respectively. For analytical tractability, we consider a Gaussian likelihood p(D ②) xx exp(-3h for
E(O; D)
[[f(x; 2 p=1
-) yyll?
(2)
where B > 0 is an inverse temperature parameter that sets the variance of the likelihood and θ = (W(4) W} [23]. We then introduce the Bayes posterior over parameters given these data:
e D) p(D O)p(e)
p(D)
we denote averages with respect to this distribution by (-). By tuning 3, one can then adjust whether the posterior dominated by the prior (3 < 1) or the likelihood (3 > 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 B oc then enforces the interpolation constraint.
2.2 The Gaussian process limit
We consider the limit of large hidden layer widths nı; n22 ......... nd-11 + oo with no. nas p. and fixed. More precisely, consider limit in which ne agn for l = I,... d 1, where ar (0, oo) and

---

- oc, studied by [3-15, 17-19, 24] and others. Importantly, we note that size of no relative to is unimportant for our results, whereas na/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 Gaussian process (GP) [3-8]. Moreover, with choice of 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(e-)) of the inner product kernel the postactivations of the final hidden layer,
k(d 1) (x,x') nz-1 W) 2/(x' (),,,,
(4)
multiplied by the prior variance of [3-8]. For a broad range of network architectures, K(S-1) can be computed recursively [5-8]. For brevity, we define the kernel matrix evaluated on the training data: [K((-1)] = K(d-1) (x,...))
3 Elementary perturbation theory for finite Bayesian neural networks
We first present our main result, which shows that the form of the leading perturbative correction to the average hidden layer kernels of 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.
3.1 Finite-width corrections to the posterior cumulants of hidden layer observables
Our main result as follows:
Conjecture Consider a BNN of the form (1), with posterior (3). Assume that this network admits well-defined GP limit as discussed in $2.2. Let O be hidden layer observable, that is, function of the hidden layer activations that is not a function of the readout weights Wa. Assume that tends in probability to a finite, deterministic limit Ox 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 k(d-1), In particular, the asymptotic expansion of the posterior mean (0) has leading terms
P
(0) = =E -nd
Σ p.A=1
lo 1o-2r 2 GgyT r-1 JAs COV (O, ((d--)) +..., pd
(5)
where I = =<d-1))) 3- Here, the cumulants of the kernels 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 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 t/. We will focus on the postactivation kernels K(0), 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 corrections are deviations of (()k()) from k(() These terms reflect corrections to the prior, and do not reflect non-trivial representation learning they are independent of the outputs. For fully-com culle-comeected networks, also known as multilaye perceptrons (MLPs), work
by Yaida [12] and by Gur-Ari and colleagues [18, 19] shows that (() K! + On1 The
3

---

((()) second type of correction is the output-dependent term that depends on cov jar 12d-1)). For deep linear MLPs MLPs with a single hidden layer, E is exactly equal to K any width (see Appendix C) [3, 12, 18], and only the covariance term contributes. More broadly, these prior works show that covw K2p,, jur, px ) (n=)) for MLPs, and that higher cumulants are of O(n [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 Seature- 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 na, 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 na/n does not tend zero with increasing n, the infinite-width behavior is not described by standard GP [8, 10]. Moreover, we note that the matrix I is invertible at any finite temperature, even when K-S-1) singular. Therefore, provided that one can extend the GP kernel by continuity to non-invertible Gza. Conjecture can be applied in the data-dense regime no < as well as the data-sparse regime no p. Furthermore, we observe that the correction depends on the outputs only through their Gram matrix Ggy- This result is intuitively sensible, since with our choice of likelihood and prior the supcte posterior is invariant under simultancous rotation of the output activations and targets. Finally, Gyy is transformed by factors of the matrix 1'_1, hence the correction depends on certain interactions between the output similarities and the GP kernel K!
3.2 High- and low-temperature limits of the leading correction
To gain some intuition for the properties of the leading finite-width corrections, we consider their high- and wwwteeppeaturee limits. These limits correspond to tuning the posterior (3) be dominated by the prior or the likelihood, respectively. At high temperatures (5 < 1), expanding r' as Neumann series (see Appendix and [27]) yields
0. TT
yy [ r- = -B + (B03)" + =((-1)) ({(a))))
(6)
Thus, at high temperatures, the outputs only influence the average kernels of Conjecture to subleading order in both width and 3, which reflects the fact that the likelihood is discounted relative to the prior in this regime. Moreover, the leading output-dependent contribution averages together Ggy and K !) hence, intuitively, there is no way to 'cancel' the GP contributions to the average kernels. We note that, at infinite temperature (3 = 0), the posterior reduces to the prior, and all finite-width corrections to the average kernels arise from the discrepancy between Ew (()) and K(9
At low temperatures (3 > 1), the behavior of r- differs depending on whether or not k<d-1) is of full rank. Assuming for simplicity that it is invertible, we have
21' Gyy 1-
[
- (d -- ") (d + +0(30) '1;
(7)
in the non non-inoertibee invertible case there are additional contributions involving projectors onto the null space of (d)1) Therefore, the leading-orde 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.
4 Learned representations in tractable network architectures
Having derived the general form of the leading perturbative Enite-width correction to the average feature kernels, 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.
4.1 Deep linear fully-connected networks
We first consider deep linear fully-connecte networks with no bias terms. Concretely, we consider a network with activations (0) "h recursively defined via (0) n
n- (0) h (e-1) with base case

---

h h()) = x. where the prior distribution of weights is [W[ Mi.d. (N For such network,
the hidden layer kernels J =\t ne ' h have deterministic limits K = m where m = oî ... of is the product of prior variances up to layer l. Higher prior cumulants of the kernels are easy to compute with the aid of Isserlis' theorem for Gaussian moments (see Appendix C) [28, 291, yielding
(K m
Guu
Σ 1=
Garl
- 1 Guz +(
Tep
where = Gaz + 1p/ (8m3) and e= l....,dd 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 more complex recurrence relation.
Thus, the leading corrections to the normalized average kernels /m3 are identical across all hidden layers up to a scalar factor that encodes the width dedth-deeeedeece 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 deep linear network. For a network with hidden layers of equal width n, we have the simple linear dependence Ee= _ (ma/ng) ) = mal/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, have
(K(9) m
Gaz +
Σ Dd
(m G GrE) On
(9)
The
in the regime in which Gaz 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 the least-norm pseudoinverse solution to the underlying underdetermined linear system X Y; we comment on this property in Appendix G.
We can gain some additional understanding of the structure of the correction by using the eigende- composition of Guz- As GIæ is by definition a real positive semidefinite matrix, it admits unitary eigendecomposition Gzz = UA with non-negative eigenvalues Apg- In this basis, the average kernel is
UUUKU A m
Σ ng
(m7² A GoUR AA) O(n
(10)
where we have defined the diagonal matrix A = elements of Ả are bounded as S Åmp S 1. Thus, the factors of T by which Ggy is conjugated have the effect of suppressing directions in the projection of Ggy onto the eigenspace of Gan with small eigenvalues. We can see that this effect will enhanced at high temperatures (3 & 1) and small scalings (m² < 1), and suppressed at low temperatures and large scalings. For this linear network, similarities are not enhanced, only suppressed. Moreover, if Græ is diagonal, then a given element of the average kernel will depend only on the corresponding element of Gys-
+ ( As 3m3A 2 0, the diagonal
We now seek numerically probe how accurately these asymptotic corrections predict learned representations in deep culnecco 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 detailed discussion of our numerical methods in Appendix Figure I, 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.
5

---

a)
b) 26
layer layer 2
Widt:: 100
[[Keo Koellr
theory
1.3
Widhh: 250
With::
Width: 500
1.0
Width: 750
0.5
Width: 100
Width: 250
107
0.0
Width:
Width: 500
Width: 750
-C.5 0.0
0.5
Width
K111
100
Figure I: Learned representations in two-hidden-layer linear fully-connected neural networks with varying widths trained via Langevin sampling on 5000 MNIST images (see Appendix for more details). (a) The Frobenius norm of the deviation of the empirical average kernel of each layer from its GP value (in this case, simply Gza) for varying widths. We see perfect match with theoretical predictions, which are shown as dashed lines. We obtain the predicted 1/n decay with increasing width and the linear scaling with the depth where the deviations for first and second layers differ by factor of 2. (b-c) Scatter plot of individual elements of the experimental (ordinate) and theoretical (abscissa) kernels for both layers. For low widths a slight deviation is visible between experiment and theory, while for larger widths the agreement is better.
4.2 Deep linear convolutional networks
To demonstrate the applicability of Conjecture to non-fully-connected BNNs, 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 Novak et al. [7] and Xiao et al. [34], we consider a network consisting of - 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 f1 where indexes the n& channels the layer and a is a spatial multi-index. The hidden layer activations are then defined through the recurrence
h (2)
Σ
h (6-1) Woj (z)
(11)
with base case f (x) ) Ti where 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 1 fu--ssddimnnsinal vector, where s is the total dimensionality of the inputs (see Appendix C for details), We fix the prior (f)
ofa)) where va is weighting factor that sets the fraction of receptive field variance location (and is thus subject the constraint La va = 1). For inputs [xpli,a and [æ,]i,a. we introduce the four-inte index hidden layer kernels
distribution of the filter elements to be Wiy. lid.
(()
1 Σ
(zw)) (O
(12)
jucab
With the given readout strategy, the two-index feature map kernel appearing in Conjecture is related to the four-index kernel of the last hidden layer by k-ld-1) 1 !. K- k((d-1) 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
[
lueaa of Σ Ue [ (a.((4+))00+))
(13)
Σm a This gives convolutional linear networks a sense spatial hierarchy that is not present in the fully-connected case: even at infinite width, the kernels include iterative spatial averaging.
with base case [a [G
In Appendix C, we derive the kernel covariances appearing Conjecture in the fully-connected case, this computation is easy to perform with the aid of Isserlis' theorem. The general result is

---

a)
Digit: 3
b)
Gyr
<)
Gyx (Tensor)
d)
G/y
333833
3333
3338
3
13333
333 33
3 3 3 333
3 993
3 3 93333
3 3
e)
leyer
V0T 150
0.00
V50 150
11Koo Korll
0.000
253
30.
00.000
w0th 662
0.000
0.109
-0.00
Width
K3
Figure 2: The MNIST image dataset and experiments for neural networks with two ID convolutional layers. (a) A 10 10 MNIST image downsized from 28 x 28 pixels.
(b) Input Gram matrix for 300 MNIST images. (c) A single (J,, V) component of the input tensor [Gralpo.ab obtained using Eq. (13). (d) The output Gram matrix. (e) The Frobenius norm of the correction the ID convolutional GP kernel is inversely proportional to the width. Here, the dashed lines are the theoretical predictions. (f) Scatter plots of individual elements of the empirical corrections to the GP kernels against the theoretical predictions for both layers show excellent agreement.
somewhat complicated, but things simplify under the assumption that readout is performed using vectorization. Then, one finds that
d-1
(K(9 pu.ab) = [K(9 Jav +
II oo
[ ap & pox 1x
Ixa +
+0n
Th
P==
P=1
p.
(14)
where we have defined pA = 3 convolutional kernel is quite similar 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-widhe kernels. In Appendix C, we show that similar simplification holds if readout is performed using global average pooling over space.
1-1 oa for brevity. Thus, the correction to the
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 2(a-d). We consider network with one-dimensional (Figure 2e and f) and two-dimensional (Figure 3) convolutional hidden layers, trained to classify 50 MNIST images (see Appendix for details our numerical methods). As shown in Figure 2(c, f) (Figure 3(a.b) for 2D convolutions), 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.
again obtain good quantitative
4.3 Networks with single nonlinear hidden layer
Finally, we would like to gain some understanding of how including nonlinearity affects the structure of learned representations. However, for nonlinear MLP, is usually not possible to analytically compute covw K<(() x<(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 studying deep nonlinear networks and gain some intuitions about how they might differ from linear BNNs. Concretely, we consider network with feature map (>; W())
7

---

a)
layer 1
b)
width 50
withh 50
layer
width 150
with 150
02
withh 400
[|Kon - Keolll
0.0
-0.4
014
10r1.
-02
-06
055
-05
10P
Width
20 3x1066122
-12
-1.
-24
-06
KC1
-07
09
B2
94
-10
Keco
K
K0L K
0.02
0.02
001
0.00
Figure 3: Learned representations in two-hidden-layer linear 2-D convolutional networks of varying channel widths. (a) The Frobenius norm of the correction the GP kernel is inversely proportional to the width. Here, the dashed lines represent theory predictions. (b) Scatter plots of individual elements of the empirical corrections to the GP kernels against the theoretical predictions for both layers show good agreement. (c) A single component (n, 7/) of the learned feature kernels in 2 2-layer CNN experiments for both convolutional layers. While the experimental kernel looks quite similar the GP (first and second columns). their difference shows the finite width corrections to the GP (last column).
10 (()x) for elementwise activation function o, where the weight matrix W(1) has prior distribution [W[1]] miid. (N (f) The only hidden layer kernel of this network is the feature map postactivation kernel Ky defined in (4), where we drop the layer index for brevity. As detailed in Appendix H, for such network have the exact expressions
[Ko] Ew ((
(15)
f1 cov Kax) [l (hr)) (r)))) - (Ko
(16)
where expectations are taken over the p-dimensional Gaussian random vector h2 which has mean zero and covariance cov( )w) of Ge Unlike for deeper nonlinear networks, here there no finite-width corrections to the prior expectations [3, 12, 18).
Though these expressions easy to define, it is not possible evaluate the four-point expectation in closed form for general Gram matrices Guæ 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
be evaluated using Isserlis' theorem (see Appendix A). However, even for a quadratic activation function Φ(x) = x², the resulting formula for the kernel will involve many elementwise matrix products, and cannot be simplified into an intuitively comprehensible form.
the input Gram matrix GG 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 activation functions that yield Eø(h) = 0 and those that yield Eœ(h) # 0. As detailed in 1, 1, if Eo(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, Ex(h) # 0, then Kx0 includes rank-1 component, and each element of the correction depends all elements of Gyy- This means that the case which Gzz is diagonal qualitatively distinct from the case in which there is only a single training input for such activation functions.
5 Learned representations in deep nonlinear networks
In the preceding section, we noted that analytical study of learned representations deep nonlinear BNNs is generally quite challenging. Here, we numerical experiments to explore whether any

---

layer
2nd layer (bottleneck)
3rd
b)
20"
3,002
|oo KGollf
3,002
Widh::
Widt: (300, 300)
(bottieneck)
3,002
30
30
Width
With: 500)
|Ken KoollF
3
Wilth 1800, s00, 500)
3,005
byer
3,004
layor bottlenecd
3003
Figure 4: 3-hidden layer neural network with ReLU activations trained via Langevin sampling on 1000 MNIST images (see Appendix 1). (a) The empirical average kernels subtracted from their corresponding GP kernels for all layers with varying widths. Labels on the y-axes indicate the widths of each layer. We observe that for networks with bottleneck layers, the deviation from K(6) is largest the bottleneck indicating representation learning; without a bottleneck deviations are considerably less (the last row). (b) Hidden layer kernel deviation from GP kernels as function of width for bottleneck networks. While the first layer shows 1 1/n scaling, the bottleneck layer and the 3rd layer deviations stay almost constant. This behavior is predicted analytically for linear networks. (c) As in (b) for networks without bottleneck. Consistent with our theory, all layers display 1/n decay.
of the intuitions gained the linear setting carry over to nonlinear networks. Concretely, we study how narrow bottlenecks affect representation learning in 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 cannot be evaluated in closed form (see $4.3). However, is straightforward simulate numerically. Consistent with the predictions of our theory for linear networks, 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 1/n with increasing width, while the deviations for the bottleneck and subsequent layers remain roughly constant. In contrast, the kernel deviations throughout network with equal-width hidden layers decay roughly as 1/n (Figure 4). These observations are qualitatively consistent with the width-dependennee of the linear network kernel (8), as well 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.
Related work
Our work closely related to several recent analytical studies of finite-width BNNs. First, Aitchison [10] argued that the flexibility afforded by finite-widh width BNNs be advantageous. He derived

---

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 finite- width prior for an MLP. However, Yaida did not attempt to study the statistics of learned features; the goal of his work was establish 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-Arr 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 perturbativelyy characterizing the internal representations that are learned upon inference.
Following the appearance of our work in preprint form, Roberts al. [36] announced an alternative derivation of the zero-temperature limit of Conjecture 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 does not require assuming 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-layee approximation to the prior could enable a fully rigorous version of Conjecture to be proved on an aebhitetuur-by-archiitee
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 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. Appendix F, we show that our result (9) for the zero-temperature kernel can be recovered as the p/n 1. 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.
7 Conclusions
In this paper, we have shown that the leading perturbative feature learning corrections to the infinite- width 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 vell-defined GP limit.
Limitations. We emphasize that our perturbative argument for Conjecture 1 is not rigorous, and that 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; pentuuruureeeive investigation of feature learning in finite BNNs will be an interesting objective for future work [9]. 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- cereetuure-b-rrcchee by-architecture basis [12, 36, 37]. Alternatively, one could attempt to treat all sufficiently sensible architectures uniformly [8, 9J. Furthermore, we have considered only one possible asymptotic regime: that in which the width is taken to infinity with 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.
10

---

Acknowledgments and Disclosure of Funding
We thank 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 's Competitive Fund for Promising Scholarship, and Google Faculty Research Award. The authors declare no conflict of interest.
References
[1] Ian Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning. MIT Press, Cambridge, MA, USA, 2016.
[2] Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. Nature, 521(7553):436-444, 2015.
[3] Radford M Neal. Priors for infinite networks. In Bayesian Learning for Neural Networks, pages 29-53. Springer, 1996.
[4] Christopher KI Williams. Computing with infinite networks. Advances in Neural Information Processing Systems, pages 295 301, 1997.
[5] Jaehoon Lee, Jascha Sohl-Dickstein, Jeffrey Pennington, Roman Novak, Sam Schoenholz, and Yasaman Bahri. Deep neural networks Gaussian processes. In International Conference Learning Representations, 2018. https htt.httts//oopereeiiee http::/://ppprrrrreeeee///ooorna http:/Eoruu11--BIA-M-----00
[6] Alexander G. de G. Matthews, Jiri Hron, Mark Rowland, Richard E. Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. International Conference on Learning Representations, 2018. URL https://openreeeee id=H1 id=H1-nGgWC-.
net/forum?
[7] Roman Novak, Lechao Xiao, Yasaman Bahri, Jachoon Lee, Greg Yang. Daniel A. Abolafia, Jef- frey Pennington, and Jascha Sohl-Dickstein. Bayesian deep convolutional networks with many channels are Gaussian processes. In International Conference on Learning Representations, 2019. URL https://openrevieu mett/forummiddB1g000000
[8] Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian pro- cess behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint 1rXiv: 2019.
[9] Greg Yang and Edward J Hu. Feature learning in infinite-width neural networks. arXiv preprint arXiv:2011.14522 2020.
[10] Laurence Aitchison. Why bigger is not always better: on finite and infinite neural networks. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning. volume 119 of Proceedings of Machine Learning Research, pages 156-164. PMLR, July 2020. URL http: http:///ooceeeenga/eeeee mlr. www.rppress/vivi. withhiisnn0aa.................a. html.
[11] Andrew Gordon Wilson and Pavel Izmailov. Bayesian deep learning and a probabilistic perspective of generalization. arXiv preprint arXiv:2002.08791. 2020.
[12] Sho Yaida. Non-Gaussian processes and neural networks at finite widths. In Jianfeng Lu and Rachel Ward, editors, Proceedings of The First Mathematical and Scientific Machine Learning Conference, volume 107 of Proceedings of Machine Learning Research, pages 165- 192, Princeton University, Princeton, NJ, USA, July 2020. PMLR. WWR:htttp:/pprceediii mlr. ppppss//117/yaaaa2///0 html.
[13] James Halverson, Anindita Maiti, and Keegan Stoner. Neural networks and quantum field theory. Machine Learning: Science and Technology, 2021.
[14] Joseph M Antognini. Finite size corrections for neural network Gaussian processes. arXiv preprint arXiv: 1908..0000000000 2019.

---

[15] Gadi Naveh, Oded Ben-David, Haim Sompolinsky, and Zohar Ringel. Predicting the outputs of finite networks trained with noisy gradients. arXiv preprint arXiv:2004.01 190,, 2020.
[16] Qianyi Li and Haim Sompolinsky. Statistical mechanics of deep linear neural networks: The backpropagating kernel renormalization. Phys. Rev. X. 11:031059, Sep 2021. doi: 10. 1103/PhysRevX.1 11.031099. URL https://link. aps 031059.
wtt.ikaaa....gg.//.../... org/ htt.11..1111/hhhh 1003///yseew..11. 11.
[17] Jacob A Zavatone-Ves Veth and Cengiz Pehlevan. Exact marginal prior distributions of finite Bayesian neural networks. In Mare' Aurelio Ranzato, Alina Beygelzimer, Perey Liang, Jenn Wortman Vaughan, and Yann Dauphin, editors, Advances in Neural Information Pro- cessing Systems, volume 34. Curran Associates, Inc., 2021. URL https: haapss///rxivv 2104. 11734.
/arxiv. org/ abs/
[18] Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from Feynman diagrams. In International Conference on Learning Representations, 2020. URL net/forum? Si-.orrr
[19] Kyle Aitken and Guy Gur-Ari. On the asymptotics of wide networks with polynomial activations. arXiv preprint arXiv:2006.0668) 2020.
[20] Florian Wenzel, Kevin Roth, Bastiaan Veeling, Jakub Swiatkowski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. How good is the Bayes posterior in deep neural networks really? In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages (0248-10255) PMLR, 13-18 Jul 2020. URL http: WRRLhttt://proceeeeee http:::///pccceeegng.eeee mlr press/
wttp./prrrss/vvsseeeeelllll v119/venze120aa html.
[21] Vincent Fortuin, Adrià Garriga-Alonso, Florian Wenzel, Gunnar Rätsch, Richard Turner, Mark van der Wilk, and Laurence Aitchison. Bayesian neural network priors revisited. arXiv preprint arXiv:2102.06571, 2021.
[22] Pavel Izmailov, Sharad Vikram, Matthew D Hoffman, and Andrew Gordon Gordon Wilson. What are Bayesian neural network posteriors really like? In Marina Meila and Tong Zhang. editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 4629 0629-4640.. PMLR, 18-24 Jul 2021. URI attp://proceeding/ mlr. www.ppees//eee hatpss/vvv/////// ht31//11191/1vvve/e// wtw.iilaaaiiiaaammmmmmmo html.
[23] David JC MacKay, A practical Bayesian framework for backpropagation networks. Neural Computation, 4(3):488 472, 1992.
[24] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. arXiv preprint 1rXiv:1806.07572, 2018.
[25] Jiri Hron, Yasaman Bahri, Roman Novak, Jeffrey Pennington, and Jascha Sohl-Dickstein. Exact posterior distributions of wide Bayesian neural networks. arXiv preprint 200.::0000.0044:: 2020.
[26] Anders Andreassen and Ethan Dyer. Asymptotics wide convolutional neural networks. arXiv preprint arXiv:2008.0860 08675, 2020.
[27] Roger A Horn and Charles R Johnson. Matrix Matrix Analysis. Cambridge University Press, 2012.
[28] Leon Isserlis. On formula for the product-moment coefficient of any order of normal frequency distribution in any number of variables. Biometrika, 12(1/2):134 139, 1918.
[29] Gian-Carlo Wick. The evaluation of the collision matrix. Physical Review, 80(2):268, 1950.
[30] Peter E Kloeden and Eekhard Platen. Stochastic differential equations. In Numerical Solution of Stochastic Differential Equations, pages 103-160. Springer, 1992.
12

---

[31] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vito,o, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, vol- ume 32. Curran Associates, Inc., 2019, URL http:: https::// ieaeiiinnnnnnuuspp...e.... cc/paper/ 2019/ ott.e///cc..22eefee2222222022220
[32] Yann LeCun, Corinna Cortes, and CJ Burges. MNIST handwritten digit database. ATT Labs [Online]. Available: http://yann.. http://yann.lecum.cmm/i www.ee...o/eee///nniiiii. 2, 2010.
[33] Roman Novak, Lechao Xiao, Jiri Hron, Jaehoon Lee, Alexander A. Alemi, Jascha Sohl- Dickstein, and Samuel S. Schoenholz. Neural tangents: Fast and easy infinite neural networks in python. In International Conference on Learning Representations, 2020. URL https: //github. com/ hotp//ooggee/eeuralloao -tangeets...
[34] Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Pen- nington. Dynamical isometry and a mean field theory of CNNs: How train 10,000-layer vanilla convolutional neural networks. In International Conference on Machine Learning, pages 5393-5402. PMLR, 2018.
[35] Devanshu Agrawal, Theodore Papamarkou, and Jacob Hinkle. Wide neural networks with bottlenecks deep Gaussian processes. Journal of Machine Learning Research, 21(175):1-66, 2020. URI.htp: / //trrmmrrr////// wwt.:/rrgpppppe///..... wttpe/rr//////0/111//// html..
[36] Daniel A Roberts, Sho Yaida, and Boris Hanin. The principles of deep learning theory. arXiv preprint arXiv:2106.10165, 2021.
[37] Boris Hanin. Random neural networks in the infinite width limit as Gaussian processes. arXiv preprint arXiv:2107.01562,, 2021.
[38] Gadi Naveh and Zohar Ringel. A self consistent theory of Gaussian processes captures feature learning effects in finite CNNs. arXiv preprint 1rXiv:2106.041100 2021.
[39] Lorenzo Noci, Gregor Bachmann, Kevin Roth, Sebastian Nowozin, and Thomas Hofmann. Precise characterization of the prior predictive distribution of deep ReLU networks. arXiv preprint 156.v:2106.066/5, 2021.
13


Reviews:
[{"review_id": "cRkja35f0Wd", "reviewer": "Reviewer_QD2J", "summary": "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. \n\n**Contributions**: \n* Updated description of Yaida [11] on the difference between finite and infinite wide neural networks \n* Feature learning interpretation \n\n**Limitations**:\n* Not rigorous \n* Works on very simple setting but no additional experiments \n* Not evident how to use the results", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Questions**: \n\n* What is the difference with the kernel used in the paper and introduced in Lee et al. “Deep neural networks as Gaussian processes”? I guess the regime when widths of layers tend to infinity sequentially can be easily addressed in the paper?\n* “As a possible complication, one could imagine a situation in which the activation distributions are ‘rotated’ during inference without changing the kernels. However, Yang and Hu [8] have shown that this cannot happen in the infinite-width limit…” Could you please precise the place in the manuscript? \n* The authors cite paper [16] in related works about the heavy-tailed nature of the prior densities over the outputs. I wonder why not cite other works that stated heavy tails like “Precise characteristics..” (the concurrent work of [16]) and Vladimirova et al. “Understanding priors…”\n\n**Quality**:\n\nI find the paper and the appendix well-written and I believe that mathematics is correct. \n\n**Significance**: \n\nThe results follow recent results of Yaida [11] and can be well-placed in the literature. However, I would like to see applications of these corrections (at least possible). \n\n**Conclusion**:\n\nThe paper has a solid theory. I recommend publication even if it misses the application part. \n\n**Minor**: \n* “Pertubative” \n* “The second class of corrections are” \n* missed comma in (A.2)\n* I would reintroduce the notations in the appendix as well, like $E$ for error function, $n$, etc \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors adequately addressed the limitations and potential negative societal impact of their work. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "30h"}, "extra_text": {}}, {"review_id": "Qrr_jOiurS8", "reviewer": "Reviewer_8cN5", "summary": "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 \tand 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 \n\nThis 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).  \n\nAuthors conjecture, non-rigorously, that these results should extend to Bayesian neural network with different architectures \n", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This is a very interesting paper, of particular interest as, despite the numerous critics of infinite-width neural networks low representational power, their learned features are not well characterized.   \n\nAuthors derive the leading finite-width corrections to the average hidden layer kernels of any fully-connected feedforward networks with linear readout and least-squares loss rigorously. I find the high- and low-temperature limits of the leading correction as well as the recovery of Aitchison implicit recurrence very compelling. \n\nThis is a first step towards understanding the features learned by Infinite width BNNs, but conjecturing that results  should extend to other architectures is an overstatement to me. Neural Networks behave in a very complex way and most related papers rely on strong assumptions on the shape of the networks to make any conclusions. The analysis of the wide “narrow bottleneck” is interesting and aims at generalizing the results but is not enough.\n\nFor better readability, Figure 1 should be updated. It is not clear that authors obtain the predicted 1/n decay with increasing width and the linear scaling with the depth, as the axis are not correctly labeled. What do the markers represent? Why do error bars don’t appear? Moreover, both figures can be made bigger. \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "QcJoe_pjGSi", "reviewer": "Reviewer_1KNh", "summary": "Work analyzing representation learning using finite-width corrections to Bayesian neural networks.  Specifically, they look at deep linear networks and networks with one nonlinear hidden layer.\n", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Work presenting insightful results on neural network theory in simple cases.  The methods are clearly developed in the Appendix.  I have no technical concerns.  There are a couple of points that I'd be interested to see discussed further (but this can easily be done in a revision):\n\nPossibly the key weakness of this work is that it uses simplistic deep linear networks, or networks with one nonlinear layer.  Of course, this allows the authors to come up with simple, interpretable results that give solid insight into the underlying processes.  But it would be worth presenting the case that these corrections are relevant in \"real world\" cases (i.e. deep nonlinear networks).\n\nLikewise, the empirical findings of Aitchison (2019) that upper-layer kernels are radically different from the infinite limit.  I'd be very interested to know whether the results of Aitchison (2019) differ from the perturbative approach, and whether that tells us anything about \"how far\" these low-order perturbations can take you from the underlying infinite limit.\n\nAppendix is very clear!  The \"Preliminary Technical Results\" section is great.\n\nGood literature review, and extensive connections to existing literature.\n\nWhile Zavatone-Veth and Pehlevan (2021) appears in the Related Work, I wasn't at all clear about the relationship between the papers (at least from the discussion there).  Would being _really_ clear about the high-level and technical relationship.  At a high-level, my guess would be that this papers looks at representation learning, while that paper focuses on the distribution of a single output variable?  But it would also be useful to sketch the technical relationship, if any!", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The idea that the paper is \"non-quantitative\" is a bit strong!", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "9otnhD9ppSs", "reviewer": "Reviewer_jCKB", "summary": "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "*Originality*: This is to my knowledge the first paper to study the leading-order pertubations to average feature kernels of general wide BNNs, and certainly I believe Conjecture 1 is original. The case of deep linear networks is closely linked to the works of Aitchison [9] and Li and Sompolinsky [15], and the authors detail this comparison in Appendix E, though the authors stress that their contributions of leading-order corrections extend these previous works.\n\n*Quality*: I believe the submission is technically sound, and written in a fair way without overstating their claims (which I commend). However, the work seems a bit light and incomplete and thus is slightly unsatisfying. Indeed, the main contribution is a non-rigorous conjecture and the analytic cases considered do not give me a great deal of insight into the inner workings of feature learning in BNNs (which is a very interesting question). \n\nSome of the questions/thoughts I have which I believe would improve the quality of this work if addressed:\n- How can you make conjecture 1 rigorous?\n- Do you get posterior concentration of the feature kernels with large p? I realise you say this is left for future work in  line 277-278, but this is quite a crucial question I think for a Bayesian perspective of wide NNs\n- Related to above, can we say anything about other moments of the finite-width corrections to feature kernels besides the posterior mean?\n- I realise that analytic tractability restricts you to deep linear and single hidden-layer, but unfortunately that means we don't see any notion of 'hierarchical' feature learning (i.e. simpler features being learned in the early layers), which intuitively is important for (B)NNs.\n- Is it wrong to think of having n_d outputs per input, as having a dataset size of $p\\times n_d$ instead of p? In which case makes the paragraph from lines 130-134 slightly unsurprising. Indeed I would expect the second term in eq 8 to scale as $O(p\\times n_d)$, but the dependence on p isn't clear?\n- Can the authors empirically demonstrate that your theory holds in more practically relevant settings. E.g. using the whole MNIST dataset as opposed to 100 points. What about CNNs with many channels?\n\n*Clarity* The paper is well written and organised. \n\n*Significance* The question of feature learning (or lack of it) in wide (B)NNs is interesting to me and I believe to the community as a whole. The results in this submission are a new angle towards understanding the interplay between the infinite-width 'lazy' GP limit and the finite-width 'feature-learning' regime, and in that sense are significant. Although see my comments in *Quality* about questions that I wish were addressed in this submission. \n\nMissed citations:\n- https://arxiv.org/abs/2106.06615 concurrent work to Zavatone-Veth and Pehlevan [6]\n- https://arxiv.org/abs/2006.10541 prove limiting posterior converges weakly to GP limit.\n\nMinor points:\n- the notation p for dataset size is confusing.\n- The conclusion focuses in on the NNGP kernel but this seems like you're selling yourself short as you also discuss earlier layers' kernels.\n- The expectation with respect to \\mathcal{W} e.g. in Eq 8 doesn't seem to be defined. I take it means the prior over weights?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors adequately address the limitations of their work in the conclusion section. They do not discuss potential negative societal impacts though I do not believe there are any.", "review_text": "", "extra_scores": {"ethical_concerns": "N/a", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper analyzes the leading-order corrections to the average feature kernels in Bayesian neural networks, and conjectures a universal form for the result. The analysis and conclusions should be of interest to theorists in the NeurIPS community, and I recommend acceptance.


Decision: Accept (Poster)