Posterior Inference on Shallow Infinitely Wide Bayesian Neural Networks under Weights with Unbounded Variance

Gaussian processes (GPs) have been studied as the infinite width limit of Bayesian neural networks with priors on network weights that have finite variance (Neal, 1996; Williams, 1996). This presents some key advantages over Bayesian neural networks with finite widths that usually require computation intensive Markov chain Monte Carlo (MCMC) posterior calculations (Neal, 1996) or variational approximations (Goodfellow et al., 2016, Chapter 19); in contrast to straightforward posterior inference and probabilistic uncertainty quantification afforded by the GP machinery (Williams and Rasmussen, 2006). In this sense, the work of Neal (1996) is foundational. The technical reason for this convergence to a GP is due to an application of the central limit theorem under the bounded second moment condition. More specifically, given an $I$ dimensional input $\mathbf{x}$ and a one-dimensional output $y(\mathbf{x})$, a $K$ layer feedforward deep neural network (DNN) with $K-1$ hidden layers is defined by the recursion:

where $z^{(1)}\equiv\mathbf{x},\;p_{1}=I,\;p_{K}=D$ and $g(\cdot)$ is a nonlinear activation function. Thus, the network repeatedly applies a linear transformation to the inputs at each layer, before passing it through a nonlinear activation function. Sometimes a nonlinear transformation is also applied to the final hidden layer to the output layer, but in this paper it is assumed the output is a linear function of the last hidden layer. Neal (1996) considers the case of a Bayesian neural network with a single hidden layer, i.e., $K=2$. So long as the hidden to output weights $w^{(2)}$ are independent and identically distributed Gaussian, or at least, have a common bounded variance given by $c/p_{2}$ for some $c>0$, and $g(\cdot)$ is bounded, an application of the classical central limit theorem shows the network converges to a GP as the number of hidden nodes $p_{2}\to\infty$.

Consider the case of a shallow, one hidden layer network, with the weights of the last layer being independent and identically distributed with symmetric $\alpha$-stable priors. Our results are derived under this setting using the following proposition of Der and Lee (2005).

Following Neal (1996), assume for the rest of the paper that the activation function $g(\cdot)$ corresponds to the sign function: $\mathrm{sign}(\xi)=1$, if $\xi>0$; $\mathrm{sign}(\xi)=-1$, if $\xi<0$; and $\mathrm{sign}(0)=0$. For $\mathbf{\xi}\in\mathbb{R}^{I}$ we define ${g(\mathbf{\xi})=\mathrm{sign}\left(b_{0}+\sum_{i=1}^{I}w_{i}\xi_{i}\right)}$, where $b_{0}$ and $w_{i}$ are i.i.d. standard Gaussian variables. The next challenge is to compute the expectation within the exponential in Equation (3). To resolve this, we break it into simpler cases. Define $\Lambda$ as the set of all possible functions $\tau:\{\mathbf{x}_{1},\dots,\mathbf{x}_{n}\}\to\{-1,+1\}$. Noting that each $\mathbf{x}_{j}$ can be mapped to two possible options: $+1$ and $-1$, indicates that there are $2^{n}$ elements in $\Lambda$. For each $\ell=1,\dots,2^{n}$, consider $\tau_{\ell}\in\Lambda$, the event ${A_{\ell}=\{\tau_{\ell}(\mathbf{x}_{j})=g(\mathbf{x}_{j})\}_{j=1}^{n}}$, and the probability ${q_{\ell}=\mathbb{P}(A_{\ell})}$. By definition $\{A_{\ell}\}_{\ell=1}^{2^{n}}$ is a set of disjoint events. Next, using the definition of the expectation of discrete disjoint events we obtain:

where the expectation is over input-to-hidden weights. A naïve enumeration sums over an exponential number of terms in $n$, and is impractical. However, details of the computation of $q_{\ell}$ and $\tau_{\ell}$ are given in Supplementary Section S.2, where we show how to reduce the enumeration over $2^{n}$ terms in Equation (4) to $L=\mathcal{O}(n^{I})$ terms using the algorithm of Goodman and Pollack (1983), by identifying only those configurations with $q_{\ell}>0$. This allows circumventing the exponential enumeration in $n$, resulting in a polynomial complexity algorithm, depending on the input dimension $I$. Although this computational complexity still appears rather high at a first glance, especially for high-dimensional inputs, for two or three-dimensional problems (e.g., in spatial or spatio-temporal models), the computation is both manageable and practical, and the complexity is similar to the usual GP regression.

Dealing with $\alpha$-stable random variables includes the difficulty that the moments of the variables are only finite up to an $\alpha$ power. Specifically, for $\alpha<2$, if $X\sim S(\alpha,\nu,\beta)$, then $\mathbb{E}[\lvert X\rvert^{r}]=\infty$ if $r\geq\alpha$, and is finite otherwise (Samorodnitsky and Taqqu, 1994, Property 1.2.16). To circumvent dealing with potentially ill-defined moments, we propose to sample from the full posterior. For fully Bayesian inference, we assign $\sigma^{2}$ a half-Cauchy prior (Gelman, 2006) and iteratively sample from the posterior predictive distribution by cycling through $(\mathbf{y}^{*},\mathbf{Q},\sigma^{2})$ in an MCMC scheme, as described in Algorithm 1, which has computational complexity of the order of $\mathcal{O}(T[(n+m)^{I}n^{2}+m^{3}])$, where $T$ is the number of MCMC simulations used. The method of Chambers et al. (1976) is used for simulating the stable variables. An implementation of our algorithms is freely available at https://github.com/loriaJ/alphastableNNet.

There are two hyper-parameters in our model: $\alpha,\nu$. We propose to select them by cross validation on a grid of $(\alpha,\nu)$, and selecting the result with smallest mean absolute error (MAE). Another possible way to select $\alpha$ is by assigning a prior. The natural choice for $\alpha$ is a uniform prior on $(0,2)$, however the update rule would need to consider the densities of the $L$ such $\alpha/2$-stable densities $p_{S^{+}}$, which would be computationally intensive as there is no closed form to this density apart from specific values of $\alpha$. A prior for $\nu$ could be included but a potential issue of identifiability emerges, similar to that identified for the Matérn kernel (Zhang, 2004). We leave these open for future research.

We compare our method against the predictions obtained from three other methods. The first two are methods for Gaussian processes that correspond to the two main approaches in GP inference: maximum likelihood with a Gaussian covariance kernel (Dancik and Dorman, 2008), and an MCMC based Bayesian procedure using the Matérn kernel (Gramacy and Taddy, 2010); the third method is a two-layer Bayesian neural network (BNN) using a single hidden layer of 100 nodes with Gaussian priors, implemented in pytorch (Paszke et al., 2019), and fitted via a variational approach. The choice of a modest number of hidden nodes is intentional, so that we are away from the infinite width GP limit, and the finite-dimensional behavior can be visualized. The respective implementations are in the R packages mlegp, tgp, and the python libraries pytorch and torchbnn. The estimates used from these methods are respectively the kriging estimate, posterior median and posterior mean. We tune our method by cross-validation over a grid of $(\alpha,\nu)$. We use point-wise posterior median as the estimate, and report the values with smallest mean absolute error (MAE) and the optimal parameters. Results on timing and additional simulations are in Supplementary Section S.4, including the posterior quantiles of $\mathbf{Q}_{\mathbf{s}}\mid\mathbf{y}$, showing the posterior is non-degenerate and stochastic. This suggests learning a non-degenerate posterior $\mathbf{Q}_{\mathbf{s}}\mid\mathbf{y}$ is possible, unlike in the GP limit where the kernel is degenerate (Aitchison et al., 2021; Yang et al., 2023). We use a data generating mechanism of the form $y=f(x)+\varepsilon$, where the $f$ is the true function. The overall summary is that when $f$ has at least one discontinuity, our method performs better at prediction than the competing methods, and when $f$ is continuous the proposed method performs just as well as the other methods. This provides empirical support that the assumption of continuity of the true function cannot be disregarded in the universal approximation property of neural networks (Hornik et al., 1989), and the adoption of infinite variance prior weights might be a crucial missing ingredient for successful posterior prediction when the truth is discontinuous.

Valuation of real estate properties in Taipei, Taiwan were collected by Yeh and Hsu (2018) in different locations. The data are available from the UCI repository¹¹1https://archive.ics.uci.edu/dataset/477/real+estate+valuation+data+set, and is a benchmark dataset. We apply our method to the spatial locations of the properties, to predict the valuations of the real estate dataset. We use 276 locations for training and 138 for testing. We compare the performance of our method to the three methods mentioned in the previous section through the mean absolute error. The results are displayed in Table 1, showing a competitive MAE under the proposed approach. Figure 5 displays the posterior predictive quantiles on the validation set, with narrower intervals under the proposed stable method in most cases.

Supplementary Sections S.5 and S.6 present results on ablation experiments and additional data sets with larger sample sizes, including recent data on S&P stock index.

We develop a novel method for posterior inference and prediction for infinite width limits of shallow (one hidden layer) BNNs under weights with infinite prior variance. While the $\alpha$-stable forward scaling limit in this case has been known in the literature (Der and Lee, 2005; Peluchetti et al., 2020), the lack of a covariance function precludes the inverse problem of feasible posterior inference and prediction, which we overcome using a conditionally Gaussian representation. There is a wealth of literature on the universal approximation property of both shallow and deep neural networks, following the pioneering work of Hornik et al. (1989), but they work under the assumption of a continuous true function. Our numerical results demonstrate that when the truth has jump discontinuities, it is possible to obtain much better results with a BNN using weights with unbounded prior variance. The fully Bayesian posterior also allows straightforward probabilistic uncertainty quantification for the infinite width scaling limit under $\alpha$-stable priors on network weights.

Several future directions could naturally follow from the current work. The most immediate is perhaps an extension to posterior inference for deep networks under stable priors, where the width of each layer simultaneously approaches infinity, and we strongly suspect this should be possible. The role of the non-degenerate posterior of $\mathbf{Q}_{\mathbf{s}}\mid\mathbf{y}$ on deep generalizations and representation learning merits a thorough investigation and suggests crucial differences from a GP limit (Aitchison et al., 2021; Yang et al., 2023). Developing analogous results under non-i.i.d. or tied weights to perform posterior inference under the scaling limits for Bayesian convolutional neural networks should also be of interest. Finally, one may of course investigate alternative activation functions, such as the hyperbolic tangent, which will lead to a different characteristic function for the scaling limit.

Our derivations rely on Equation 5.4.6 of Uchaikin and Zolotarev (1999), which states that for $\alpha_{0}\in(0,1)$ and for all positive $\lambda$ one has: ${\exp(-\lambda^{\alpha_{0}})=\int_{0}^{\infty}\exp(-\lambda t)p_{S^{+}}(t)dt},$ where $p_{S^{+}}$ is the density function of a positive $\alpha_{0}$-stable random variable. Using $\lambda=\nu z^{2}$ and $\alpha_{0}=\alpha/2$ and the fact that $z^{2}=\lvert z\rvert^{2}$, we obtain for $\alpha\in(0,2)$ that:

where $p_{S^{+}}$ is the density function of a positive $\alpha/2$-stable random variable and $z\in\mathbb{R}$. By Equation (4) of Der and Lee (2005), and the characteristic function of independent normally distributed error terms with a common variance $\sigma^{2}$, we have that:

where the third equality follows by using Equation (7), and in the last equality we define $\mathbf{M}_{\ell}$ as the matrix with ones in the diagonal and with the $(i,j)$th entry given by $\tau_{\ell}(\mathbf{x}_{i})\tau_{\ell}(\mathbf{x}_{j}),\;i\neq j$. Next, using the fact that the densities are over the independent variables $\{s_{\ell}\}_{\ell=1}^{L}$ we bring the product inside the integrals and employ the property of the exponential to obtain:

using on the second line the definition of $\mathbf{Q}_{\mathbf{s}}$. The required density is now obtained by the use of the inverse Fourier transform on the characteristic function:

where the second line follows by the derived expression for the characteristic function, and the third line follows by Fubini’s theorem since all the integrals are real and finite. We recognize that the term $\exp(-(1/2)\mathbf{t}^{T}\mathbf{Q}_{\mathbf{s}}\mathbf{t})$ corresponds to a multivariate Gaussian density with covariance matrix $\mathbf{Q}^{-1}_{\mathbf{s}}$, though it is lacking the usual determinant term. We obtain the density using the characteristic function of Gaussian variables to finally obtain the result:

In using $\mathbf{Q}^{-1}_{\mathbf{s}}$ freely, we assumed through the previous steps that $\mathbf{Q}_{\mathbf{s}}$ is positive-definite. We proceed to prove this fact. Note that $\mathbf{Q}_{\mathbf{s}}$ is obtained from the sum of $L$ rank-one matrices and a diagonal matrix, where each of the rank-one matrices is $q_{\ell}^{2/\alpha}s_{\ell}\nu\boldsymbol{\tau}_{\ell}\boldsymbol{\tau}^{T}_{\ell}$. Let $\mathbf{w}\in\mathbb{R}^{n}\backslash\{0\}$. Then:

implying that $\mathbf{Q}_{\mathbf{s}}$ is positive-definite with probability 1.

The Supplementary Material contains technical details and numerical results in pdf. Computer code is freely available at: https://github.com/loriaJ/alphastableNNet.

Bhadra was supported by U.S. National Science Foundation Grant DMS-2014371.

Posterior Inference on Shallow Infinitely Wide Bayesian Neural Networks under Weights with Unbounded Variance
(Supplementary Material)

One of the most important properties of stable random variables is the closure property (Property 1.2.1, Samorodnitsky and Taqqu, 1994), which states that if $X_{i}\sim S(\alpha,\nu_{i},\beta_{i})$ independently for $i=1,2$, then $X_{1}+X_{2}\sim S(\alpha,\xi,\gamma)$, where $\gamma=(\beta_{1}\nu_{1}^{\alpha}+\beta_{2}\nu_{2}^{\alpha})/(\nu_{1}^{\alpha}+\nu_{2}^{\alpha})$, and $\xi=(\nu_{1}^{\alpha}+\nu_{2}^{\alpha})^{1/\alpha}$. This means that the sum of two $\alpha$-stable variables is again $\alpha$-stable. This is a generalization of the well known property of the closure under convolutions of Cauchy ($\alpha=1$) and Gaussian ($\alpha=2$) random variables. In terms of moments, Samorodnitsky and Taqqu (1994, Property 1.2.16) indicate that for $X\sim S(\alpha,\beta,\nu)$ with $\alpha\in(0,2)$, we have $\mathbb{E}[\lvert X\rvert^{r}]=\infty$ for $r\geq\alpha$ and $\mathbb{E}[\lvert X\rvert^{r}]$ is finite for $0<r<\alpha$. Specifically, this implies that $\alpha$-stable random variables have infinite variance, when $\alpha<2$.

The property of closure under convolutions is easily generalized to the sum of a sequence of i.i.d. $\alpha$-stable variables, which gives rise to a convergence in a non-Gaussian domain, for $\alpha<2$. Formally, the generalized central limit theorem (Gnedenko and Kolmogorov, 1954) proves that for i.i.d. scaled random variables with infinite variance the convergence is no longer to a Gaussian random variable. Rather the convergence is to an $\alpha$-stable random variable. A statement of the theorem is below.

Finally, the Laplace transforms of positive $\alpha$-stable random variables exist. For $\alpha<1$ and $X\sim S(\alpha,1,1)$ the Laplace transform is given by:

for $\lambda>0$.

Since the size of $\Lambda$ is $2^{n}$, indicating an exponential complexity of naïve enumeration, we further simplify Equation (4). When the input points are arranged in a way that $\tau_{\ell}$ is not possible, then $q_{\ell}$ must be zero. We can identify the elements in $\Lambda$ with positive probabilities by considering arbitrary values of $b_{0},w_{1},\dots,w_{I}$. The corresponding $\tau$ function is determined by: ${\tau(\mathbf{x}_{j})=\mathrm{sign}(b_{0}+w_{1}x_{j1}+\cdots+w_{I}x_{jI})}$, for each $j$. This corresponds to labeling with $+1$ the points above a hyperplane, and with $-1$ the points that lie below the hyperplane. When $I=1$, without loss of generality, let $x_{1}<\cdots<x_{n}$. In this case $\Lambda$ corresponds to the possible changes in sign that can occur between the input variables, which is equal to $n$. Specifically, the sign change can occur before $x_{1}$, between $x_{1}$ and $x_{2}$, …, between $x_{n-1}$ and $x_{n}$, and after $x_{n}$. A similar argument is made in Example 2.1.1 of Der and Lee (2005), suggesting the possibility of considering more than one dimensions.

For $I>1$, Harding (1967) studies the possible partitions of $n$ points in $\mathbb{R}^{I}$ by an $(I-1)$-dimensional hyperplane—which corresponds to our problem, and determines that for points in general configuration there are $O(n^{I})$ partitions. Goodman and Pollack (1983) give an explicit algorithm for finding the elements of $\Lambda$ that have non-zero probabilities in any possible configuration. We summarize their algorithm for $I=2$ as Algorithm S.1 in Supplementary Section S.3. This algorithm runs in a computational time of order $n^{I}\log(n)$, which is reasonable for moderate $I$. For the rest of the article, we denote the cardinality of elements in $\Lambda$ that have positive probability by $L$, with the understanding that $L$ will depend on the input vectors $\mathbf{x}_{i}$ that are used and their dimension. This solves the issue of computing deterministically the values of $\tau_{\ell}$ that have positive probability after integrating through input-to-hidden weights.

Next we compute the probability $q_{\ell}$ for the determined $\tau_{\ell}$. For $I=1$, the $q_{\ell}$s correspond to probabilities obtained from a Cauchy cumulative density function, which we state explicitly in Supplementary Section S.2.1. For general dimension of the input $I>1$, the value of $q_{\ell}$ is given by $\mathbb{P}(\mathbf{Z}^{(\tau_{\ell})}>0)$, where the $n$-dimensional Gaussian vector $\mathbf{Z}^{(\tau_{\ell})}$ has $i$-th entry given by $\tau_{\ell}(\mathbf{x}_{i})(b_{0}+\sum_{j=1}^{I}w_{j}x_{ij})$. This implies that $\mathbf{Z}^{(\tau_{\ell})}\sim\mathcal{N}_{n}(0,\boldsymbol{\Sigma}^{(\tau_{\ell})})$, where the (possibly singular) variance matrix is ${\Sigma^{(\tau_{\ell})}_{i,j}=\tau_{\ell}(\mathbf{x}_{i})\tau_{\ell}(\mathbf{x}_{j})(1+\sum_{k=1}^{n}x_{ik}x_{jk})}$. This means we can compute $q_{\ell}=\mathbb{P}(\mathbf{Z}^{(\tau_{\ell})}>0)$ using for example the R package mvtnorm, which implements the method of Genz and Bretz (2002) for evaluating multivariate Gaussian probabilities. Since the $q_{\ell}$s require independent procedures to be computed, this is easily parallelized after obtaining the partitions.