\section{Introduction}


Bayesian neural networks (BNNs) provide principled notions of uncertainty in deep learning, but they have not been widely adopted in practice as many properties of this model and its inference remain poorly understood, particularly in the overparameterized, nearly nonparametric regime where modern deep learning take often places. One tool for understanding the behavior of BNNs with a large number of parameters is to take the limit as the number of hidden units (i.e., \emph{width}) goes to infinity. In this case, the prior predictive distribution of a single-layer BNN converges in distribution to the \textit{NNGP}, a Gaussian process with the \emph{neural network kernel} that depends on the prior and architecture of the network \citep{neal_1996}. Extensions of this result exist also for deep networks  \citep{matthews_2018} and for BNN posteriors  \citep{hron_2020}. 

In contrast, asymptotic properties of popular variational approximations of BNN posteriors have not been extensively studied.  In the finite width regime, variational posterior approximations are known to have deficiencies.  For example, the commonly used mean-field variational posterior, which ignores correlations between parameters, exhibits various pathologies including poor uncertainty estimates between data rich regions \citep{foong_2019_2}. Moreover, recent works have noted the tendency of mean-field variational BNN posteriors to underfit even with large network architectures \citep{dusenberry2020efficient}. It is therefore natural to ask if the deficiencies in mean-field variational approximations of finite width BNN posteriors persist in approximations of infinite width BNN posteriors.

In this paper, we show that the answer, unfortunately, is \emph{yes}. For single-layer Bayesian neural networks used for univariate regression, we prove a surprising result: the posterior predictive mean under mean-field variational inference converges to zero (assuming the observed outcomes are centered) as the number of hidden units tends to infinity. That is, the mean-field variational posterior predictive mean completely ignores the data, unlike the true posterior predictive of an infinite width BNN. Furthermore, we provide insight on the cause of this failure --- we show that this results from the over-regularizing effect of the KL divergence term (forcing the posterior to match the zero-mean prior) in the variational inference objective. 


\section{Background}

\paragraph{Bayesian neural networks (BNNs)}

A single-layer feedforward neural network with $K$ hidden units used for univariate regression is given by:
\begin{equation}
	f(x, \theta) = \sum_{k=1}^K w^{(2)}_k \psi(w^{(1)T}_k x + b^{(1)}_k )
	\label{eq:bnn}
\end{equation}
where $\psi$ is a nonlinear activation function, $x\in \mathbb{R}^D$ is an input, $w^{(1)}_k \in \mathbb{R}^D$ and $b^{(1)}_k \in \mathbb{R}$ are input-layer weight and bias parameters, respectively, and $w^{(2)}_k \in \mathbb{R}$ is an output-layer weight parameter. For simplicity, we assume the mean has been subtracted from the observed outcomes so that we can omit an output-layer bias parameter. We let $\theta \in \mathbb{R}^{K(D+2)}$ denote the collection of all model parameters. Given independent and identically distributed observations $\mathcal{D} = \{(x_n, y_n)\}_{n=1}^N$, we assume a Gaussian likelihood function $\mathcal{L}(\theta) = \prod_n \mathcal{N}(y_n \mid f(x_n, \theta), \sigma^2_\text{noise})$ and infer the parameters $\theta$ by maximizing the likelihood.

A \textit{Bayesian neural network} (BNN) places a prior distribution $p(\theta)$ on the parameters, typically a factorized Gaussian:
\begin{align}
	w^{(1)}_k&\sim \mathcal{N}(0, \sigma_{w^{(1)}}^2 I_D)
	\label{eq:prior1}
	\\
	b^{(1)}_k &\sim \mathcal{N}(0, \sigma_{b^{(1)}}^2)
	\label{eq:prior2}
	\\
	w^{(2)}_k &\sim \mathcal{N}(0, \sigma_{w^{(2)}}^2),
	\label{eq:prior3}
\end{align}
where $I_D$ is a $D\times D$ identity matrix and $k\in \{1,\dotsc,K\}$.


\paragraph{Convergence to Gaussian processes} 
Under this prior, the covariance between the function $f$ evaluated at two inputs $x$ and $x'$ is given by:
\begin{align*}
	K  \sigma_{w^{(2)}}^2 \mathbb{E}[\psi(w^{(1)T}_k x + b^{(1)}_k) \psi(w^{(1)T}_k x' + b^{(1)}_k)]
\end{align*}
Notice by scaling the output-layer prior variance parameter $\sigma_{w^{(2)}}^2$ inversely with the network width $K$ (i.e., setting $\sigma_{w^{(2)}}^2 = \tilde{\sigma}_{w^{(2)}}^2/K$ for some $\tilde{\sigma}_{w^{(2)}}^2$), the prior covariance is the same for any width. Then, as the number of hidden units tends to infinity, application of the central limit theorem reveals that for any set of inputs $x$, the prior predictive distribution over the function output $f(x)$ approaches a multivariate Gaussian with the covariance given above. This  the definition of a \textit{Gaussian process} (GP). In this case, it is called the neural network Gaussian process (NNGP), since the covariance function of the multivariate Gaussian distribution over $f(x)$ is induced by the neural network.

Note that while scaling the output-layer prior variance parameter by $1/K$ is useful for analyzing the theoretical properties of a BNN, it is an important practical consideration, too. Otherwise the prior predictive variance could become too large as the width increases. 

\paragraph{Variational inference}
Unfortunately, the posterior distribution of a finite-width BNN is not available in closed-form and Markov chain Monte Carlo (MCMC) methods are too slow for all but the smallest networks. Instead, it is common to find the closest distribution $q_\phi$ in KL divergence to the posterior by maximizing a lower bound on the marginal likelihood called the evidence lower bound (ELBO):
\begin{equation}
	\label{eq:elbo}
	\text{ELBO}(\phi) = \mathbb{E}_{\theta\sim q_\phi}[\log \mathcal{L}(\theta)] - \mathbb{K}\mathbb{L}[q_\phi || p(\theta)].
\end{equation}
The first term in the ELBO is the expected log likelihood, which measures how well the model fits the data, and the second term is the Kullback-Leibler (KL) divergence regularization, which measure how close $q_\phi$ is to the prior $p(\theta)$. 

The \textit{variational} distribution $q_\phi$ is parameterized by a set of variational parameters $\phi$. A common choice for $q_\phi$ is a product of independent (i.e., \textit{mean-field}) Gaussian distributions, one distribution for each parameter $\theta_i$ in the model:
\begin{equation}
	q_\phi(\theta) = \prod_{i=1}^{|\theta|} \mathcal{N}(\theta_i \mid \mu_i, \sigma^2_i).
	\label{eq:q}
\end{equation}
Here, $\phi = \{(\mu_i, \sigma^2_i)\}$ are the variational parameters. We call variational inference using Equation \ref{eq:q} \textit{mean-field variational inference} (MFVI).

Since both the prior and variational distribution are Gaussian, the KL divergence can be calculated in closed-form. For example, the KL divergence between the variational distribution of all $K$ output-layer weights $w^{(2)}$ and a $\mathcal{N}(0, \frac{1}{K}I_K )$ prior is 
\begin{align}
	\frac{1}{2} \sum_{k=1}^{K} \left[K \mu_k^2 + K\sigma_k^ 2 -1 - \log K \sigma_k^2 \right],
	\label{eq:kl}
\end{align}
where $\{\mu_k\}$ and $\{\sigma^2_k\}$ are the variational parameters. Notice Equation \ref{eq:kl} acts like L2 regularization of the mean parameters, which will be key to our proof of Theorem \ref{th:1}.


\begin{figure}[H]
	\centering
	\includegraphics[width=.47\textwidth]{figures/figure1.png}
	\vspace{-.3cm}
	\caption{Prior and posterior predictive distributions for mean-field variational BNNs of different widths compared to the NNGP, to which the true posterior of the BNN converges. For a large width, the mean-field variational BNN ignores the data, unlike the NNGP. The shaded regions constitute $\pm 1$ standard deviation around the means (solid lines). All estimates are based on 1000 function samples (a few of which are drawn faintly).}
	\label{fig:posteriors}
\end{figure}

\section{Main result} \label{sec:main}

Figure \ref{fig:posteriors} illustrates our main observation. The left column compares the prior predictive distributions of BNNs of various widths (top panel) to the Gaussian process prior to which these BNNs converge (i.e., the NNGP, bottom panel). As expected from the scaling in the prior, the prior predictive mean and variance is the same between all finite and infinite width models. However, the right column reveals significant differences in the posterior predictive distributions. While the NNGP posterior exhibits a posterior mean that models the data and a posterior variance that expands outside of the data, the MFVI BNN posterior predictives completely ignore the data as the width increases. We use the implementation provided by \cite{novak_2019} to compute the NNGP.


Like other BNN pathologies observed in practice, it is not immediately obvious whether the behavior is due to the variational objective itself or poor optimization of the ELBO. However, in Theorem \ref{th:1}, we show that the underfitting observed in wide BNNs is due to the choice of the variational objective itself (i.e., the combination of the choice of the variational family, the prior and the divergence measure in variational inference). Specifically, we prove that the MFVI posterior mean converges to zero as the width approaches infinity. Although we assume an error function activation function, we empirically demonstrate that the same conclusions likely hold other activations (we provide results for ReLU and tanh activations in Section \ref{sec:experiments}).

\begin{theorem}
	\label{th:1}
	Consider a mean-field variational BNN of width $K$ as described by equations \ref{eq:bnn} and \ref{eq:prior1}-\ref{eq:prior3} and assume an error function activation $\phi=\mathrm{erf}(z):=\int_{0}^z 2/\sqrt{\pi} \exp(-t^2)~dt$. Suppose $\hat{\phi}_K$ maximizes the ELBO given in Equation \ref{eq:elbo}. Then for any dataset $\mathcal{D}$ and input $x^*$, the variational posterior predictive mean $\mathbb{E}_{\theta \sim q_{\hat{\phi}_K}}[f(x^*,\theta)]$ converges to zero as the number of hidden units $K\to\infty$.
\end{theorem}


To understand why Theorem \ref{th:1} holds, think of the single-layer neural network in Equation \ref{eq:bnn} as a linear model using $K$ nonlinear basis functions $\psi(w^{(1)T}_k x + b^{(1)}_k )$. Upon observing data, as the number of basis functions $K$ increases, the prior encourages each of them to be weighted less in the posterior, because of the $1/K$ scaling of the prior variance of the output-layer weights $w^{(2)}$. This tradeoff is what allows the prior predictive variance to converge as $K\to\infty$. 

Unfortunately, because of the KL regularization of input-layer parameters $w^{(1)}$ and $b^{(1)}$ in the ELBO, basis functions cannot freely be added to the model without further penalty. Recall that the number of observations is fixed, so there is little improvement to the expected log likelihood term in the ELBO once the neural network has enough hidden units to fit the data. This limits the degree to which the variational distribution over the input-layer parameters can differ from the prior and, in the limit of $K$, prevents the model from fitting the data. 

In contrast, a random features model (e.g., Random Fourier Features \cite{rahimi_2007}), where the input-layer parameters are drawn from a fixed distribution at initialization, has no such issue. Additional basis functions are drawn from the same distribution for all $K$ and can be added to the model without penalty. 

The formal proof can be found in the Appendix, but we give a brief sketch here. The proof of Theorem \ref{th:1} follows in roughly three steps. 
\begin{itemize}
\item \textbf{Step 1: evaluate the ELBO at the prior}
Notice the ELBO evaluated at the prior does not depend on the width $K$. This provides a lower bound on the ELBO that holds for networks of any width.

\item \textbf{Step 2: bound the KL regularization}
Because the ELBO cannot be lower than the bound, the KL divergence cannot be higher than the bound. In particular, this constrains the L2 norm of the optimal variational mean parameters (see Equation \ref{eq:kl}).

\item \textbf{Step 3: bound the posterior mean} 
Using the constraint on variational mean parameters, application a few basic inequalities to the posterior mean reveals convergence to zero. 
\end{itemize}

While Theorem \ref{th:1} shows that the variational posterior predictive mean converges to the prior predictive mean, we provide empirical evidence that the variational posterior predictive variance converges also to the prior predictive variance. In particular, Figure \ref{fig:varational_params} in the supplementary material shows that the variational posterior variances approaches the prior variances and Figure \ref{fig:convergence} shows that the variational predictive variance becomes more similar to the prior predictive variance as we increase network width. We thus conjecture that a stronger statement is true: that the mean-field variational posterior distribution converges to the prior. A proof of this is current work. Finally, although Theorem \ref{th:1} applies to single hidden layer networks, understanding the asymptotic properties of deep networks is also work in progress. 



\begin{figure}[H]
	\centering
	\includegraphics[width=.48\textwidth]{figures/figure3.pdf}
	\vspace{-.8cm}
	\caption{Across datasets, the mean-field variational posterior of a BNN gets closer to the prior as the width increases. \textit{Top:} Euclidean distance between the prior and posterior predictive mean. \textit{Bottom:} Euclidean distance between the prior and posterior predictive variance. We use 50 test points spaced uniformly over a grid centered roughly around the training data.}
	\label{fig:convergence}
\end{figure}

\begin{figure}[H]
	\centering
	\includegraphics[width=.45\textwidth]{figures/figure2.png}
	\vspace{-0.5cm}
	\caption{Comparison of the posterior predictives of MFVI BNNs using tanh and ReLU activation functions for small (left panel) and large (right panel) widths. Regardless of the activation, wide MFVI BNNs ignore the data.}
	\label{fig:relu}
\end{figure}

\section{Experiments}
\label{sec:experiments}

We begin by analyzing the rate of convergence of the posterior predictive mean of the mean-field variational BNN to zero for different datasets. The datasets are the ``2-points'' dataset shown in Figures \ref{fig:posteriors} and \ref{fig:relu}, a synthetic dataset of noisy observations of a sine wave ($N=20$), and a real dataset of winning Olympic marathon paces over the period 1896 to 2012 ($N=27$).\footnote{Available in the open-source \texttt{pods} Python package \url{https://github.com/sods/ods}.} We $z$-score standardize the inputs and outputs of all datasets. 

For each dataset and widths ranging from $K=125$ to $K=32,000$, the top and bottom panels, respectively, of Figure \ref{fig:convergence} show the Euclidean distance of the posterior predictive mean  and variance under mean-field variational inference to the prior predictive mean and variance. The top panel illustrates the conclusions of Theorem \ref{th:1}: as the width increases, the posterior predictive mean converges to zero. In the bottom panel panel, we see the posterior predictive variance getting closer to the prior, though it is still far away. We will investigate this behavior in future work. 


Next we analyze the impact of different activation functions. Although Theorem \ref{th:1} assumes an erf activation function, Figure \ref{fig:relu} provides empirical evidence that the results of Theorem \ref{th:1} hold for tanh and ReLU activations as well. 




\subsection{Implementation details}
We use a prior variance of 2 for all parameters (i.e., $\sigma_{w^{(1)}}^2=2$,  $\sigma_{b^{(1)}}^2=2$, and $\tilde{\sigma}_{w^{(2)}}^2=2$) and a prior observational noise $\sigma^2_{\text{noise}}=.01$. However, we implement the impact of the prior distribution by scaling the parameters by their prior variance in the forward pass and then using a $\mathcal{N}(0,1)$ distribution as the prior to evaluate the KL divergence term in the ELBO. That is, for any parameter $\theta_0 \in \theta$, where $p(\theta_0) = \mathcal{N}(0, \sigma_{\theta_0}^2)$ is its prior distribution, we replace $\theta_0$ with $\sigma_{\theta_0} \theta_0 $ in the evaluation of $f(x, \theta)$ and  use a $\mathcal{N}(0,1)$ as the prior. In particular, this means that instead of scaling the output-layer prior variance $\sigma_{w^{(2)}}^2$ by $1/K$, we scale the function output by $1/K$ and keep $\sigma_{w^{(2)}}^2$  unscaled. This ``neural tangent kernel'' scaling yields the same prior predictive distribution while enabling a constant learning rate to be used for training networks of different widths \cite{lee_2019}.  

We initialize the variational mean and variance parameters from a normal-inverse-gamma family. Specifically, if $q(\theta_0) = \mathcal{N}(\mu, \sigma^2)$ is the variational distribution corresponding to the parameter $\theta_0$, we randomly initialize $\mu\sim\mathcal{N}(0,1)$ and $\sigma^2 \sim \mathcal{I}\mathcal{G}(\nu+1, \nu)$. It follows from the laws of total expectation and variance that $\mathbb{E}[\theta_0] = 0$ and $\mathbb{V}[\theta_0] = 2$. The hyperparameter $\nu$ controls the concentration of $\sigma^2$ around its initial mean of one (i.e., $\mathbb{E}[\sigma^2]=1$ and $\mathbb{V}[\sigma]=1/(\nu-1)$). We set $\nu$ to the width of the network. We have experimented with other initializations, including very small variances and  mean parameters pretrained to maximize the log likelihood, with little impact on the overall results. 







\section{Discussion}


\citet{trippe_2018} discusses \textit{over-pruning}, which is the phenomenon whereby many of the variational distributions over the output-layer weights concentrate around zero. This is undesireable behavior because the amount of over-pruning increases with the degree of over-parameterization and because over-pruning degrades performance --- simpler models that do not permit pruning often perform better. As in our work, the explanation for over-pruning centers around the tension between the likelihood term and the KL divergence term in the ELBO. To reduce the KL divergence, hidden units can be pruned from the model. Relatedly, we prove that as the number hidden units tends to infinity, the KL divergence over-regulates the model by pulling each of the output-layer weights towards zero while limiting the overall probability mass assigned to the input-layer weights (see earlier discussion in Section \ref{sec:main} regarding the intuition of Theorem 1). 

Our work also provides theoretical insight into \textit{cold posteriors}, which is the empirical phenomenon that downweighting the importance of the KL divergence in the ELBO yields better model performance \cite{wenzel_2020}. It is possible this practice serves to undo the over-regularization of the KL divergence that we investigate. 

A common theme in both of these pathologies is that the effect becomes more prominent as the width of the network increases. Yet, the phenomenon of \textit{double descent} shows that it is in this over-parameterized regime where neural networks perform best \cite{belkin_2019}. Therefore, it is critical to understand the properties of wide variational BNNs --- which we prove are considerably different from the true posterior in the mean-field case --- if they are to be adopted in practice. 



