Asymptotics of representation learning in finite Bayesian neural networks

1 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 of their hidden layer representations is only theoretically well-understood in certain infinite-width limits , in which these representations can not 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 . 2 Preliminaries . 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 . 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 : 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 . 2.2 The Gaussian process limit . 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ν ) . 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 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 . 3.1 Finite-width corrections to the posterior cumulants of hidden layer observables . 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 ) ∞ .

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.