Statistical Foundation of Variational Bayes Neural Networks

Bayesian neural networks (BNNs) have been comprehensively studied in the works of Bishop [1997], Neal [1992], Lampinen and Vehtari [2001], etc. More recent developments which establish the efficacy of BNNs can be found in the works of Sun et al. [2017], Mullachery et al. [2018], Hubin et al. [2018], Liang et al. [2018], Javid et al. [2020] and the references therein. The theoretical foundation of BNN by Lee [2000] widens the scope to a broader community. However, with the age of big data applications, the conventional Bayesian approach is computationally inefficient. Thus the alternative computational approaches, such as variational Bayes (VB) become popular among machine learning and applied researchers. Although, there have been many works on the algorithm development for VB in recent years, the theoretical advancement on estimation accuracy is rather limited. This article provides statistical validity of neural networks models with variational inference along with some theory-driven practical guidelines for implementation.

In this article, we mainly focus on feed-forward neural networks with a single hidden layer of inputs and a logistic activation function. Let the number of inputs be denoted by $p$ and the number of hidden nodes by $k_{n}$ where the number of nodes is allowed to increase as a function of $n$. The true regression function, $E(Y|X=\boldsymbol{x})=f_{0}(\boldsymbol{x})$ is modeled as a neural network of the form

where $\psi(u)=1/(1+\exp(-u))$ is the logistic activation function. With a Gaussian-prior on each of the parameters, Lee [2000] establishes the posterior consistency of neural networks under the simple setup where the scale parameter $\sigma=V(Y|X=\boldsymbol{x})$ is fixed at 1. The results in Lee [2000] mainly exploit Barron et al. [1999], a fundamental contribution that laid down the framework for posterior consistency in non parametric regression settings. In this paper, we closely mimick the regression model of Lee [2000] by assuming $y=f_{0}(\boldsymbol{x})+\xi$ where $f_{0}(\boldsymbol{x})$ is the true regression function and $\xi$ follows $N(0,\sigma^{2})$.

The joint posterior distribution of a neural network model is generally evaluated by popular Markov Chain Monte Carlo (MCMC) sampling techniques, like Gibbs sampling, Metropolis Hastings, etc. (see, Neal [1996], Lee [2004], and Ghosh et al. [2004] for more details). Despite the versatility and popularity of MCMC based approach, the Bayesian estimation suffers from computational costs, scalability, time constraints along with other implementation issues such as choice of proposal densities and generating sample paths. Variational Bayes emerged as an important alternative to overcome the drawbacks of the MCMC implementation (see Blei et al. [2017]). Many recent works have discussed the application of variational inference to Bayesian neural networks e.g., Logsdon et al. [2009], Graves [2011], Carbonetto and Stephens [2012], Blundell et al. [2015], Sun et al. [2019]. Although, there is a plethora of literature implementing variational inference for neural networks, the theoretical properties of the variational posterior in BNNs remain relatively unexplored and this limits the use of this powerful computational tool beyond the machine learning community.

Some of the previous works that focused on theoretical properties of variational posterior include the frequentist consistency of variational inference in parametric models in the presence of latent variables (see Wang and Blei [2019]). Optimal risk bounds for mean-field variational Bayes for Gaussian mixture (GM) and Latent Dirichlet allocation (LDA) models have been discussed in Pati et al. [2017]. The work of Yang et al. [2017] propose $\alpha$ variational inference Bayes risk for GM and LDA models. A more recent work Zhang and Gao [2017] discusses the variational posterior consistency rates in Gaussian sequence models, infinite exponential families and piece-wise constant models. In order to evaluate the validity of a posterior in non-parametric models, one must establish its consistency and rates of contraction. To the best of our knowledge, the problem of posterior consistency, has not been studied in the context of variational Bayes neural network models.

Our contribution: Our theoretical development of posterior consistency, an essential property in nonparametric Bayesian Statistics, provides confidence in using the variational Bayes neural networks model across the disciplines. Our theoretical results help to assess the estimation accuracy for a given training sample and model complexity. Specifically, we establish the conditions needed for the variational posterior consistency of the feedforward neural networks. We establish that a simple Gaussian mean-field approximation is good enough to achieve consistency for the variational posterior. In this direction, we show that $\varepsilon$- Hellinger neighborhood of the true density function receives close to 1 probability under the variational posterior. For the true posterior density ( Lee [2000]), the posterior probability of an $\varepsilon$- Hellinger neighborhood grows at the rate $1-e^{-\epsilon n^{\delta}}$. In contrast, we show for the variational posterior this rate becomes $1-\epsilon/n^{\delta}$. The reason for this difference is explained by two folds: (1) first, the KL-distance between the variational posterior and the true posterior does not grow at a rate greater than $n^{1-\delta}$ for some $0\leq\delta<1$, (2) second, the posterior probability of $\varepsilon$- Hellinger neighborhood grows at the rate $1-e^{-\epsilon n^{\delta}}$, thus, the variational posterior probability must grow at the rate $1-\epsilon/n^{\delta}$, otherwise the rate of growth of the KL-distance cannot be controlled. We also give the conditions on the approximating neural network and the rate of growth in the number of nodes needed to ensure that the variational posterior achieves consistency. As a last contribtuion, we show that the VB estimator of the regression function converges to the true regression function.

Further, our investigation shows that although the variational posterior(VP) is asymptotically consistent, posterior probability of $\varepsilon-$Hellinger neighborhoods does not converge to 1 as fast as the true posterior. In addition, one requires that the absolute value of the parameters in the approximating neural network function grow at a controlled rate (less than $n^{1-\delta}$ for some $0\leq\delta<1$), a condition not needed in dealing with MCMC based implementation. When the absolute value of the parameters grow as a polynomial function of $n$ ($O(n^{v}),v>1$), one can choose a flatter prior (a prior whose variance increases with $n$) in order to guarantee VP consistency.

VP consistency has been established irrespective of whether $\sigma$ is known or unknown and the differences in practice have been discussed. It has been shown that one must guard against using Gaussian distributions as a variational family for $\sigma$. Since the KL-distance between variational posterior and true posterior must be controlled, one must ensure that quantities like $E(\log X)$ and $E(1/X^{2})$ are defined under the variational distribution of $\sigma$. We thereby discuss two sets of variational family on $\sigma$, (1) an inverse gamma-distribution, (2) a normal distribution on the log-transformed $\sigma$. While the second approach may seem intuitively appealing if one were to use fully Gaussian variational families, it comes with a drawback. Indeed, under the reparametrized $\sigma$, the variational posterior is consistent if the rate of growth in the number of nodes is slower than under the original parameter models. However, a smaller growth in the number of nodes makes it more and more difficult to find an approximating neural network which converges fast enough to the true function.

The outline of the paper is as follows. In Section 2, we present the notation and the terminology of consistency for variational posterior. In Section 3, we present the consistency results when the scale parameter is known. In Section 4, we present the consistency of an unknown scale parameter under two sets of variational families. In Section 5, we show that the Bayes estimates obtained from the variational posterior converge to the true regression function and scale parameter. Finally, Section 5 ends with a discussion and conclusions from our current work.

In this paper, we assume that the true regression function $f_{0}$ can be approximated by a neural network

where $k_{n}$, the number of nodes increases as a function of $n$, while $p$, the number of covariates is fixed. Thus, the total number of parameters grow at the same rate as the number of nodes, i.e. $K(n)=1+k_{n}(p+1)\sim k_{n}$.

Suppose there exists a neural network $f_{\boldsymbol{\theta}_{0n}}(\boldsymbol{x})=\beta_{00}+\sum_{j=1}^{k_{n}}\beta_{j0}\psi(\gamma_{j0}^{\top}\boldsymbol{x})$ such that

Note that if $f_{0}$ is a neural network function itself, then (A1) holds trivially for all $0\leq\delta<1$ irrespective of the choice of $k_{n}$. Theorem 2 of Siegel and Xu [2019] showed that with $k_{n}=n$, $\delta$ can be chosen between $0\leq\delta<1/2$. Mimicking the steps of Theorem 2, Siegel and Xu [2019], it can be shown that with $k_{n}=n^{a},a>1/2$, $\delta$ can be chosen anywhere in the range $0\leq\delta<a-1/2$. For a given choice of $k_{n}$, whether (A1) holds or not depends on the entropy of the true function. Assumptions of similar form can also be found in Shen [1997] (see conditions C and $C^{\prime}$) and Shen et al. [2019] (see condition C3).

Note that the condition (A1) characterizes the rate at which a neural network function approaches to the true function. The next set of conditions characterize the rate at which the coefficients of the approximating neural network solution grow. Suppose, one of the following two conditions hold:

Note that condition (A2) ensures that sum of squares of the coefficients grow at a rate slower than $n$. White [1990] proved consistency properties of feed forward neural networks with $\sum_{i=1}^{K(n)}|\theta_{i0n}|=o(n^{1/4})$ which implies $\sum_{i=1}^{K(n)}|\theta_{i0n}|^{2}\leq(\sum_{i=1}^{K(n)}|\theta_{i0n}|)^{2}=o(n^{1/2})$, i.e. $0\leq\delta<1/2$. Blei et al. [2017] studied the consistency properties for parametric models wherein one requires the assumption $-\log p(\theta_{0})$ be bounded (see Relations (44) and (53) in Blei et al. [2017]). With a normal prior of the form $p(\boldsymbol{\theta}_{n})\propto\exp(-\sum_{i=1}^{K(n)}\theta_{in}^{2})$, the same condition reduces to $\sum_{i=1}^{K(n)}\theta_{i0n}^{2}$ bounded at a suitable rate. Indeed, condition (A2) guarantees that the rate of growth KL-distance between the true and the variational posterior is well controlled.

Condition (A3) is a relaxed version of (A2), where the sum of squares of the coefficients is allowed to grow at a rate in polynomial in $n$. A standard prior independent of $n$ might fail to guarantee convergence. We thereby assume a flatter prior whose variance increases with $n$ in order to allow for consistency through variational bayes. Note that if $f_{0}$ is a neural network function itself, conditions (A2) and (A3) hold trivially.

In this section, we begin with the simple model where the scale parameter $\sigma_{0}$ is known. For a simple Gaussian mean field family as in (13), we establish that variational posterior is consistent as long as assumptions (A1), (A2) and (A3) hold. We also discuss, how the rates contrast with those in Lee [2000] which established the posterior consistency of the true posterior.

In the following two theorems for two class of priors, we establish the posterior consistency of $\pi^{*}$, i.e. the variational posterior concentrates in $\varepsilon-$ small Hellinger neighborhoods of the true density $l_{0}$. Note that, assumptions (A2) and (A3) impose a restriction on the rate of growth of the sum of squares of the coefficients of the approximating neural network solution. With (A2), we show that a standard normal prior on all the parameters works. However, under the more weaker assumption (A3), a normal prior whose variance increases with $n$ is needed. Additionally, we show that for the variational posterior to achieve consistency, the number of parameters or equivalenty the number of nodes $k_{n}$ need to grow in a controlled fashion.

Note that conditions (16) and (17) agree with those assumed in Theorem 1 of Lee [2000]. Since $\pi^{*}(\mathcal{V}^{c}_{\varepsilon})=o_{P_{0}}(n^{-\delta})$, the variational posterior is consistent with $\delta$ as small as 0. Indeed $\delta=0$ imposes the least restriction on the convergence rate and coefficient growth rate of the true function (see assumptions (A1) and (A2)). As $\delta$ grows, restrictions on the approximating neural function increase but that guarantees faster convergence of the variational posterior. Expanding upon the Bayesian posterior consistency established in Lee [2000], one can show that $\pi(\mathcal{V}_{\varepsilon}^{c}|\boldsymbol{y}_{n},\boldsymbol{X}_{n})\leq o_{P_{0}^{n}}(e^{-n^{\delta}})$ for any $0\leq\delta<1$ (see Relation (88) in Lee [2000]). Thus, probability of $\varepsilon-$ Hellinger neighborhood grows at the rate $1-\epsilon(1/n)^{\delta}$ for variational posterior in contrast to that of $1-\epsilon(e^{-n})^{\delta}$ for true posterior. For parametric models, the rate of growth of the variational posterior was found to be $1-\epsilon(1/n)$ (see second equation 2 on page 38 of Blei et al. [2017]). Note that the consistency of true posterior requires no assumptions on the approximating neural network function whereas for the variational posterior, both assumptions (A1) and (A2) must be satisfied to guarantee convergence.

Observe that the consistency rate in Theorem 3.2 agrees to the one in Theorem 3.1. In order to prove both theorems 3.1 and 3.2, a crucial step is to show that $d_{KL}(\pi^{*}(.),\pi(.|\boldsymbol{y}_{n},\boldsymbol{X}_{n}))=o_{P_{0}^{n}}(n^{1-\delta})$. In order to show this, we show that $d_{KL}(q(.),\pi(.|\boldsymbol{y}_{n},\boldsymbol{X}_{n}))=o_{P_{0}}(n^{1-\delta})$ for some $q\in\mathcal{Q}_{n}$. Indeed this choice of $q$ varies in order to adjust for changing nature of the prior from (17) to (18) (see statements (1) and (2) in Lemma 7.9).

We next present the proof of Theorems 3.1 and 3.2. The first crucial step of the proof is to establish that the $d_{KL}(\pi^{*}(.),\pi(.|\boldsymbol{y}_{n},\boldsymbol{X}_{n}))$ is bounded below by a quantity which is determined by the rate of consistency of the true posterior (see the quantities $A_{n}$ and $B_{n}$ in the proof below). The second crucial step towards the proof is to show $d_{KL}(\pi^{*}(.),\pi(.|\boldsymbol{y}_{n},\boldsymbol{X}_{n}))$ is bounded above at a rate which can be greater than the rate of its lower bound iff the variation posterior is consistent.

In the above proof we have assumed $\pi^{*}(\mathcal{V}_{\varepsilon})>0$, $\pi^{*}(\mathcal{V}_{\varepsilon}^{c})>0$. If $\pi^{*}(\mathcal{V}_{\varepsilon}^{c})=0$, there is nothing to prove. If $\pi^{*}(\mathcal{V}_{\varepsilon})=0$, then following the steps of the proof, we will get $\varepsilon^{2}=o_{P_{0}}(n^{-\delta})$ which is a contradiction.

The second quantity $B_{n}$ is controlled by the rate at which the prior gives mass to shrinking KL neighborhoods of the true density $l_{0}$. Indeed, the quantity $B_{n}$ appears again when computing bounds on $d_{KL}(q(.),\pi(.|\boldsymbol{y}_{n},\boldsymbol{X}_{n})$ for some $q\in\mathcal{Q}_{n}$ (see $③$ in Proposition 7.19). If $\delta=0$, $B_{n}$ can be controlled even without assumptions (A1) and (A2). However, if $\delta>0$, assumptions (A1) and (A2) are needed in order to guarantee that the $B_{n}$ grows at a rate less than $n^{1-\delta}$.

The last quantity, $d_{KL}(\pi^{*}(.),\pi(.|\boldsymbol{y}_{n},\boldsymbol{X}_{n}))$ is controlled at a rate less than $n^{1-\delta}$ by showing that there exists a $q\in\mathcal{Q}_{n}$ (see (62) and (65)) such that $d_{KL}(\pi^{*}(.),\pi(.|\boldsymbol{y}_{n},\boldsymbol{X}_{n}))=o_{P_{0}^{n}}(n^{1-\delta})$. Both assumptions (A1) and (A2) play an important role in guaranteeing that such a $q$ does exist.

In this section, we now assume that the scale parameter $\sigma$ is unknown. In this case, our approximating variational family is slightly different from (14). Whereas, we still assume a mean field Gaussian family on $\boldsymbol{\theta}_{n}$, our approximating family for $\sigma$ cannot be Gaussian. An important criterion to guarantee the consistency of variational posterior is to ensure $\int d_{KL}(l_{0},l_{\boldsymbol{\omega}_{n}})q(\boldsymbol{\omega}_{n})d\boldsymbol{\omega}_{n}$ is well bounded (see Lemma 7.11). When $\sigma$ is unknown, $d_{KL}(l_{0},l_{\boldsymbol{\omega}_{n}})$ involves terms like $\log\sigma$ and $1/\sigma^{2}$ both of whose integrals are undefined under a normally distributed $q$. We thereby adopt two versions of $q$ for $\sigma$, firstly an inverse gamma distribution of $\sigma$ and secondly a normal distribution on the log transformed $\sigma$ (see Sections 4.1 and 4.2 respectively). Both the transforms have their respective advantage in terms of determining the rate of consistency of the variational posterior. In this section, we work only with assumption (A2). We can handle (A3) in a way exactly similar to Section 3.