Bayesian inference with finitely wide neural networks

Neal in his seminal work [1] pointed out that a shallow but infinitely wide random neural network is a Gaussian process (GP) [2] in statistical sense. Subsequent work [3, 4] in interpreting neural network with specific nonlinear activation units as kernel machines was also inspired by such idea. More recent reports [5, 6] further claimed the equivalence between GP and deep neural networks when each hidden layer in latter is of infinite width. Consequently, machine learning practitioners can perform Bayesian inference by treating deep and wide neural network as a GP, and exploit the analytic marginal and conditional properties of multivariate Gaussian distribution. Otherwise, one needs to employ gradient-based learning and bootstrap sampling for obtaining predictive distribution [7].

In reality, all neural networks have finite width. Therefore, the deviation from Gaussianity requires further quantitative account as practitioners may wonder the corrections to the predictive mean and variance in, for example, a regression task. Yaida [8] and colleagues [9] proposed a perturbative approach for computing the multivariate cumulants by direct application of Wick’s contraction theorem. Moreover, the fourth cumulants are shown to be nonzero, scaled by the sum of reciprocal widths, $1/N_{1}+1/N_{2}+\cdots+1/N_{L-1}$, and signaling non-Gaussian aspect of the random processes representing finite-width deep neural networks with $L$ hidden layers [8]. A quartic energy functional for a fixed set of network outputs is formulated field-theoretically with which the corrections to the posterior mean and variance due to the weak non-Gaussianity are obtained [8].

While Yaida’s approach is appealing from a field theory perspective (see also [10, 11]), the loss of elegant marginal property due to the presence of fourth power term in exponent is critical for analytic inference with finite-width networks. An alternative thinking is to modify the multivariate Gaussian distribution so that the new distribution can match the higher cumulants associated with the networks. In this paper, we shall use the multivariate Edgeworth series expansion [12] to construct the non-Gaussian distribution for the network’s outputs. In particular, we find that the differential representation of Edgeworth expansion greatly facilitates the derivation of marginal and conditional properties of the non-Gaussian distribution. Three main results are reported in this paper. First, the marginal property is intact, and the corrections to conditional mean and variance are derived. Second, with observed data, the non-Gaussian posterior distribution associated with an unobserved output is derived. Third, we derive the marginal covariance [13] of a bottlenecked deep neural network [14], which represents deep Gaussian Process [15] in weight space. It is worthwhile to note that some of the hidden layers in the bottlenecked network are narrow and strong non-Gaussianity may be induced.

The paper has the following organization. In Sec. II we begin by reviewing the computational structure of a shallow cosine network, its equivalence to GP with a Gaussian kernel, and the emergence of non-Gaussianity due to finite width. The shallow network with random parameters in Bayesian setting is a non-Gaussian prior over function. In Sec. III the non-Gaussian prior is represented as a differential representation of Edgeworth expansion around Gaussian multivariate distribution, and its marginal and conditional properties are established. Application of the non-Gaussian prior in Bayesian regression task is discussed in Sec. IV. Finally, Sec. V is devoted to investigating the combined effect of nonlinear activation, depth and finite width on the non-Gaussian prior for a deep bottlenecked cosine network, which is followed by a discussion in Sec. VI.

Let us start with the discussion of a random shallow feed forward network with cosine activation. In the infinite width limit, the network is statistically equivalent to a GP with Gaussian kernel [16]. The goal is to see the emergence of nonzero fourth cumulant when the width is finite. Consider a single output network with $N$ activation units, the real function value $z_{\alpha}\in\mathbb{R}$ with greek subscript is an indexed random variable in association with its input ${\bf x}_{\alpha}\in\mathbb{R}^{d}$. Explicitly, the output-input of the network has the following relation,

where the weight variables $w$’s are sampled from the Gaussian distribution $\mathcal{N}(0,1)$, the scaling variables $\Omega$’s sampled from $\mathcal{N}(0,I_{d})$, and the phase variables $\phi$’s sampled from uniform distribution $\mathcal{U}([0,2\pi])$. The normalization factors $\sqrt{2/N}$ and $1/\sqrt{d}$ above follow the parameterization used in [17]. Because $w$ is zero-mean, the first relevant statistical moment is the covariance, $k({\bf x}_{\alpha},{\bf x}_{\beta})$, namely the expectation of product of two function values at two inputs,

In above, the independence between the random variables $w$’s is used to arrive at the second equality. The third equality is due to the fact that Fourier transformation of Gaussian is Gaussian and the vanishing average of cosine with uniformly random phase. The fourth moment can be computed in similar manners, but one shall notice that the product of two cosine terms containing the random phase can generate a nonzero contribution. Here, we focus on the fourth cumulant tensor as the signature of non-Gaussianity,

Hence, the fourth cumulant tensor receives a contribution proportional to the reciprocal width, $1/N$, and it is symmetric with respect to permutation of indices.

The fourth cumulant tensor in Eq. (3) signifies that the distribution over the network outputs is non-Gaussian. In other words, the random parameters $w$’s, $\Omega$’s, and $\phi$’s from a prior distribution induce the non-Gaussian prior function distribution due to the finite width. The first interesting consequence of the nonzero fourth cumulant is that the prior distribution over a single output, say $z_{\alpha}$, receives a correction term deviating from Gaussian [18],

The variance $\sigma_{0}^{2}:=k(z_{\alpha},z_{\alpha})$ and the new distribution $q$ has matched moments, i.e. $\mathbb{E}_{q}[1]=1$, $\mathbb{E}_{q}[z_{\alpha}^{2}]=\sigma_{0}^{2}$, and $\mathbb{E}_{q}[z_{\alpha}^{4}]=3\sigma_{0}^{4}+V_{\alpha\alpha\alpha\alpha}$. Although the work [18] derived the above non-Gaussian distribution from a renormalization group perspective, such expansion with Hermite polynomials around Gaussian distribution has been known as Edgeworth expansion in statistics literature [19, 20]. When the width $N\rightarrow\infty$, the fourth cumulant $V$ vanishes and for notational convenience we denote the limiting distribution by $q_{\infty}$, which is Gaussian.

Having established the marginal and conditional properties of the non-Gaussian distribution in Eq. (5), we now continue investigating the posterior distribution over the unseen function value $z_{*}$ associated with input ${\bf x}_{*}$ when the noisy observations $\mathcal{D}=\{({\bf x}_{1},y_{1}),({\bf x}_{2},y_{2}),\cdots\}$ are known. According to Bayes’s rule, the objective distribution is the posterior distribution defined as,

with $\Lambda$ denoting a diagonal matrix representing the noise in Gaussian likelihood. When the $q$’s in the numerator and denominator are both Gaussian, which corresponds to the infinitely wide network, the limiting distribution reads

In above first equality, the first and second Gaussian distributions on the right hand side represent the conditional density $q_{\infty}(z_{*}|{\bf z})$ and the posterior $q_{\infty}({\bf z}|\mathcal{D})$, respectively.

As for the case of finite-width, the involved Gaussian likelihood can be dealt with by continuing application of the properties of adjoint differential operator as well as derivative of independent Gaussians. The details of derivation of evidence $q(\mathcal{D})$ can be seen in Appendix C, and the expression in terms of derivatives with respect to the observed $y$’s reads,

Despite its simple form, the evidence term can not be further be manipulated like in the infinite-width case. The details of derivation of the posterior distribution for output $z_{*}$ at the unseen input,

can also be found in Appendix C. Here, the notations can be understood as follows: the hatted indices, $\hat{i}$ for instance, additionally include the symbol $*$ associated with test function value $z_{*}$, and the derivative refers to $\partial_{*}:=\partial_{z_{*}}$ (with respect to function value) and $\partial_{i}:=\partial_{y_{i}}$ (with respect to the observed value).

We conclude this Section with considering the same simple regression example but with noisy observation in $y_{2}$ at input ${\bf x}_{2}$, and we wish to predict the unobserved function value $z_{1}$. To stress the role of noise parameter $\sigma_{n}$, we only show the predictive mean here,

with $\alpha_{n}=1/\sqrt{1+\sigma_{n}^{2}}$. The correction term vanishes as the noise parameter does, and is odd with respect to the sign change of $y_{2}$.

Up to this point, we have focused on the emergence of non-Gaussianity in prior function distribution in wide and shallow network, as well as how the non-Gaussian prior affects the inference in a regression task. Our approach can be extended to the deep and wide neural networks as the corresponding prior distributions approach GP [5, 6]. However, such wide-width assumption leading to weak non-Gaussianity is not valid for the bottlenecked networks in which the hidden layers are alternatively wide and narrow [22, 14].

Let us consider a two-layer bottlenecked feed forward network modeling the hierarchical mapping ${\bf x}\in\mathbb{R}^{d}\mapsto{\bf z}^{(1)}\in\mathbb{R}^{H}\mapsto z^{(2)}\in\mathbb{R}$. The hierarchy consists of the following computations,

where the parameters $w$’s, $\Omega$’s, and $\phi$’s share the same prior distribution with their counterparts in the shallow network. The widths of hidden layers, $N_{1}$ and $N_{2}$, are large but the hidden output dimension $H$ is not. In the limit $N_{1,2}\rightarrow\infty$ while $H$ remains finite, the two-layer cosine bottlenecked network corresponds to the prior of deep Gaussian process [15, 23, 24] with Gaussian kernel, which is a flexible and expressive function prior due to its compositional nature. As its name suggests, the conditional prior $z^{(2)}|{\bf z}^{(1)}$ is a GP and so is each component $z_{i}^{(1)}|{\bf x}$ in the hidden output. However, the marginal distribution for $z^{(2)}|{\bf x}$ is non-Gaussian as [13] showed that the fourth cumulant is positive, i.e., a heavy-tailed distribution [25].

Here, we are interested in tracking how the small parameters, $1/N_{1}$ and $1/N_{2}$, and the bottleneck parameter $1/H$ enters the second moment. For a deep and finitely wide linear network, the prior is non-Gaussian [8, 26] but the second moment does not receive correction due to the finite width [8, 9]. Thus, the effects of nonlinear activation on higher statistical moments are interesting (also see a recent work in [27]). For the deep cosine network in Eq. (20), we can first consider the wide limit, $N_{1,2}\rightarrow\infty$, but the bottleneck width $H$ remains finite. Following the result in Eq. (2), the covariance for the deep model can be computed as followings,

Note that the decomposition $\mathcal{N}(z_{1},z_{2}|0,K)=\mathcal{N}(z_{1}+z_{2}|0,\sigma_{+}^{2})\mathcal{N}(z_{1}-z_{2}|0,\sigma_{-}^{2})$ [28] has been used in derivation with the variances $\sigma_{\pm}^{2}=2(K_{11}\pm K_{12})$. We want to stress that the above kernel is exact for any $H$. For a very narrow case, i.e., $H=1$, the above result coincides with that in [13], while in the very wide bottleneck case, $H\gg 1$, it can be shown that

where the covariance $k$ in shallow network is given in (2). In addition, the exponential of shallow kernel in above corresponds to $H\rightarrow\infty$, which is the same as the deep kernel in [29]. Indeed, as suggested in [9], the appearance of reciprocal width, $1/H$, in the correction term is in association with the weak non-Gaussianity even though the other widths $N_{1,2}$ are infinite.

Next, we shall compute the covariance for the case where all widths are finite. In fact, the outer width $N_{2}$ does not enter the kernel and we only need large but finite $N_{1}$. The computation is similar and one only needs to replace the limiting joint distribution $q_{\infty}(z^{(1)}_{\alpha},z^{(1)}_{\beta})$ with the non-Gaussian $q$. The details of derivation using the property derivative of Gaussian can be seen in Appendix D. The kernel of the two-layer bottlenecked cosine network with $N_{1,2}<\infty$ reads,

with the correction due to finite $N_{1}$,

with $\sigma_{-}^{2}=2(1-k)$ is used for easing the notation. Again, the above result is for general $H$. For the wide bottleneck limit, one can show that these small parameters enter the kernel with the correction of $O(\frac{1}{H})$ followed by $O(\frac{1}{N_{1}H})$. Hence, the deep model as a function prior serves as a GP with random kernel whose mean value converges to Eq. (22) when all widths $N_{1,2}$ and $H$ approach infinity. However, the role of $H$ is different from $N_{1}$ from our perturbative analysis, and $1/N_{1}$ does not appear alone in the small parameter. The influence of finite width on the kernel of deep and nonlinear and convolutional models is also studied in [30, 31, 32].

The fourth moment can be computed in a similar manner. The closed form for $N_{1,2}\rightarrow\infty$ can be found in [13]. For the finite width case, the fourth cumulant $V_{\alpha\beta\gamma\delta}^{(2)}$ becomes the $H$-th power of the permutational symmetrization of $\frac{1}{N_{2}}\mathbb{E}_{q}\{\exp[-(z^{(1)}_{\alpha}+z^{(1)}_{\beta}-z^{(1)}_{\gamma}-z^{(1)}_{\delta})^{2}/2H)\}$, the small parameters of which shall consist of $1/N_{2}$, $1/(N_{2}H)$, and $1/(N_{2}HN_{1})$. The analysis of scaling with respect to the reciprocal width is more complex than the deep linear network.

In essence, the finite width in random neural networks induces nonzero variance of kernel since the fourth cumulant is not zero. From function space perspective, the shallow networks with finitely large width can be regarded as a GP but the kernel itself is a random variable too, so the learned data representation is a distribution over kernels. Therefore, neural network with very narrow layers may not have enough expressive power for learning while the network with very wide layers is not flexible enough because it is equivalent to GP learning with one fixed kernel. It was recently suggested in Ref. 33 through studying the partition function [34, 35] of finite deep network that Student-t process [21] may suitably represent finite-width network, which offers an alternative description than this paper. The reports in [30, 24] observed the degradation of performance for Bayesian deep neural network when expanding the widths while the study in [36] suggested otherwise.

Taking the wide limit with the perturbative approach is appealing because of the analytic and elegant expressions for inference, which may shed light on future investigation of alternative algorithm for classification [37], study of average test error in regression [38, 39] and classification [40] tasks using finite-width networks. In this paper, the conditional statistics and perturbed posterior distribution over the unseen output using the non-Gaussian prior are obtained with help of the differential representation of multivariate Edgeworth expansion. In parallel with the equivalence between GP prior and the random wide neural network, the investigation of neural tangent kernel [17, 41, 42] is important for understanding the learning dynamics during the optimization and the relation with Bayesian neural network learning [43, 44].

The research is supported by the Dean Office of School of Arts and Science at Rutgers University Newark. Correspondences with Jacob Zavatone-Veth and Pietro Rotondo are acknowledged.