On weight and variance uncertainty in neural networks for regression tasks

Bayesian Neural Networks (BNNs) have been introduced and comprehensively discussed by many authors (among others see [14, 2]). BNNs are suitable for modeling uncertainty by considering values of the parameters that might not be learned by the available data. This is achieved by randomizing unknown parameters, such as weights and biases while incorporating prior knowledge. Such an approach naturally regularizes the model and helps prevent overfitting, a common challenge in neural networks (NNs). The BNNs can also be considered as a suitable and more reliable alternative to ensemble learning methods [15], including bagging and boosting of the neural networks.

The primary objective in the BNNs is to estimate the posterior distribution of the parameters. However, obtaining the exact posterior is often infeasible due to the integrals’ intractability. Thus, it is necessary to use approximation techniques for the posterior density. One common approach is the Markov Chain Monte Carlo (MCMC), including the Metropolis-Hastings algorithm [10], which generates samples from the posterior distribution by constructing a Markov chain, typically of first orders.

Despite its effectiveness, flexibility, and strong theoretical foundations, MCMC has significant limitations in high-dimensional spaces, including slow convergence and high computational costs. Variational Bayes (VB) has emerged as a faster and more scalable approximation technique. First introduced by [11] and [19], VB seeks to approximate the intractable posterior distribution by identifying a simpler, parametric probability distribution that closely resembles the true posterior. The central idea in VB is to recast the problem of posterior inference as an optimization task. Specifically, VB minimizes the Kullback-Leibler (KL) divergence [13] between the approximate distribution and the true posterior [3]. This involves iteratively updating the hyper-parameters of the approximate distribution to improve its fit to the posterior while maintaining computational efficiency [1].

The VB method has limitations related to the conjugacy of the full conditionals of parameters. In such cases, models that go beyond the conjugacy of the exponential family have been proposed [20]. Black Box Variational Inference (BBVI) [16] eliminates the need for conditional conjugacy by introducing a flexible framework for approximating posterior distributions. It leverages stochastic optimization techniques to estimate gradients and employs variance reduction strategies to enhance computational efficiency and stability. Bayes by Backprop (weight uncertainty), introduced by Blundell et al. [4], is another approach for performing Bayesian inference in NNs. This method leverages stochastic optimization and employs the re-parametrization trick to enable efficient sampling and gradient-based optimization of posterior distributions, for a NN weights. The model is used for classification and regression tasks, by considering suitable multi-nomial and Gaussian likelihood for the target variable, respectively.

Another approach for modeling uncertainty in weights of a NN is considered in [7]. They modeled the variational approximation of the weights as the multiplication of a fixed hyperparameter (which might be optimized) by a bernoulli random varaible for modeling the droupout mechanism. There are two main differences between the method considered in [7] and that in [4]. The first difference is that no statistical distribution is considered for modelling uncertainty in weights around the hyperparameter in the method proposed in [7]. The second difference is the approach for modeling the dropout mechanishm, since the dropout mechinism was modeled by considering the slab-and-spike prior for the weights of the NN in [4].

For the regression task, the variance of the Gaussian likelihood is assumed fixed in [4] and is determined by cross-validation. In this paper, we consider the problem of variance uncertainty in the model proposed by [4] to investigate its effectiveness on the prediction performance of the BNNs for regression tasks. Throughout a simulation study from a nonlinear function approximation problem, it is shown that the Bayes by Backprop with variance uncertainty performs better than the model with fixed variance. We have also considered the riboflavin data set, which is a genetic high dimensional data set to evaluate the performance of the proposed model, compared to the other methods, by applying PCA-BNN and the dropout-BNN methods. The dropout-BNN is applied by considering the spike-and-slab prior for the weights of the BNN. As the results presented in Section 6 shows, our proposed method outperforms other approaches in terms of MSPE and coverage probabilities for the regression curve, as well as MSPE for the riboflavin dataset.

The remainder of the paper is organized as follows. Section 2 reviews the foundations of the BNNs. The VB method for posterior approximation is introduced in Section 3. Section 4 describes the Bayes by Backprop method proposed by [4], for the regression task with fixed variance. In Section 5, we propose the variance uncertainty for the BNNs. The experimental results and the evaluation of the proposed model is considered in Section 6. Some concluding remarks are given in Section 7. The codes for this study is available online on github (soon).

In this section, our goal is to describe BNNs and their challenges in practical applications. Before examining BNNs, we briefly describe NNs from a statistical perspective. A NN aims to approximate the unknown function $y=\phi(x)$ where $x$ and $y$ are the input and the output vectors, respectively. From a statistical point of view, given the sample points $x,y$, the objective function of NN is the log-likelihood function of weights and biases, $W$, and extra parameters $\rho$, of the model, $\log L(W,\rho|x,y)$ which is aimed to be maximized to learn the parameters of the network

The likelihood function $L(W,\rho\mid x,y)$ depends on the specific problem being addressed:

• •

  In regression tasks, the likelihood function is typically modeled by a Gaussian density of $y$ with a mean provided by the network given the input data set $x$ and an unknown variance $\sigma^{2}$, and the negative log-likelihood corresponds to the sum of squared errors (SSE).

• •

  In classification problems, the likelihood function is often modeled as a categorical distribution, and the negative log-likelihood corresponds to the cross-entropy loss.

In a BNN, a prior distribution $p(W,\rho)$ is considered for the parameters of the model [14], and our goal is to find a posterior distribution of all parameters $p(W,\rho\mid x,y)$. Using the Bayes’ formula, the posterior density is obtained as

where $p(x,y)$ is the marginal distribution of the data, given by

The point estimation of the model parameters might be computed by the maximum a posteriori (MAP) estimator as follows

The posterior predictive distribution function, used to make predictions and to determine prediction intervals for new data, is defined as follows (for more details, see e.g., [14, 2])

In practice, we face intractable integrals to solve $p(x,y)$; therefore, we resort to approximation methods. In the following, we introduce VB as a well-known approximation technique.

The VB is a method for finding an approximate distribution $q(\theta)$ of the posterior distribution $p(\theta|x)$, by minimizing the Kullback-Leibler divergence ${\rm KL}\left[q(\theta)||p(\theta|x)\right]$ as a measure of closeness [18]. In this section, we briefly review the basic elements of the VB method.

Let $x$ be a vector of observed data, and $\theta$ be a parameter vector with joint distribution $p(x,\theta)$. In the Bayesian inference framework, the inference about $\theta$ is done based on the posterior distribution $p(\theta|x)=p(x,\theta)/p(x)$ where $p(x)=\int p(x,\theta)d\theta$.

In the VB method, we assume the parameter vector $\theta$ is divided into $M$ partitions $\{\theta_{1},\ldots,\theta_{M}\}$ and we want to approximate $p(\theta|x)$ by

where $q_{j}(.)$ is the variational posterior over the vector of posterior parameters $\theta_{j}$. Hence, each $\theta_{j}$ can be a vector of parameters. The best VB approximation $q^{\ast}(\cdot)$ is then obtained as

where

Let $\eta=\{\eta_{1},\ldots,\eta_{M}\}$ be the partitions of hyper-parameter such that each $\eta_{j}$ be the variational parameter of $\theta_{j}$. In the fixed-form-VB (FFVB), we consider an assumed density function for each $q(\theta_{j}|\eta_{j})$, with unknown $\eta_{j}$, and we aim to determine the hyper-parameters $\eta$ such that ${\rm KL}\left[q(\theta|\eta)||p(\theta|x)\right]$ is minimized.

It is clear from (1) that minimizing KL is equivalent to maximizing the evidence lower bound (ELBO) with respect to $\eta$, defined as

The ELBO is equal to $\log p(x)$ when the KL divergence is zero, meaning a perfect fit. The closer the value of ${\rm ELBO}[q(\theta)]$ to $\log p(x)$ (greater values), the fit is better. Maximizing the ELBO is equivalent to minimizing its negative. It follows that

where $\eta^{*}$ represents the optimal parameter estimate for maximizing the ELBO.

We define a loss function to be minimized. Assume that $\epsilon$ is a random variable with a probability density function $q(\epsilon)$, and let $f$ be a function. In this case, as shown in [4], if

then

Equation (3) shows that if we can find a transformation for $\theta$ such that it depends on $\epsilon$ and satisfies the condition (2), then samples of $\epsilon$ can be used to generate samples of $\theta$. This approach focuses on finding $\eta^{*}$, as defined above. Therefore, we define

Note that the expectation of $f(\theta,\eta)$ is the negative of the ELBO and the gradient of $f(\theta,\eta)$ provides an unbiased estimate of the gradient defined in equation (3). In addition, increasing the number of samples of $\epsilon$ reduces the variance of the estimated gradient. This method utilizes equation (4) as an objective function to determine $\eta^{*}$.

In this section, we consider regression problems with constant variance. As we mentioned in section 3, we can use equation (4) for the optimization problem as the KL error and we are interested in minimizing it.

Following the approach described in [4], the variance of the likelihood function in a regression problem can be fixed to a constant value. In regression problems, assuming independency among the output features (i.e., the values of the vector $y_{i}$), we have

where $I_{q}$ is the identity matrix, $x_{i}\in\mathbb{R}^{p}$ represents the $i$-th input variable, $y_{i}\in\mathbb{R}^{q}$ denotes the $i$-th output variable, $i=1,...,n$, and $\sigma^{2}_{0}$ is a fixed known variance. Moreover, $\phi(x_{i};W)$ represents the output of a NN with an arbitrary number of layers. In this case, we have the conditional distribution

If $(x_{i},y_{i})$ is the $i$th observation, then the likelihood function, assuming independent observations, is given by

Thus, [4] considered the objective function

with

and $\eta=(\mu_{w},\rho_{w})$, seeking the optimum value of $\eta$. We denote this method by VBNET-FIXED throughout the remaining of the paper. However, this approach fails to effectively capture uncertainty in the variance parameter. In the forthcoming section, we will discuss how to introduce uncertainty in the variance by using parametrization, so that the model can be generalized by considering the posterior distribution over the variance parameter.

In this section, we suggest to use a variational posterior distribution over weights, biases, and also the variance of the likelihood function. Our proposal can be considered as a generalization of the previous model introduced in section 4. In our approach, we assume that the posterior parameters are represented by $\theta=(W,S)$, where $W$ encompasses the weights and biases of the NN, and $S$ denotes a parameter related to the variance of the likelihood function. Additionally, we define the variational parameters as $\eta=(\sigma_{w}^{2},\mu_{w},\sigma_{L}^{2},\mu_{L})$, where $\sigma_{w}^{2}$ and $\mu_{w}$ pertain to the variational distribution over $W$, and $\sigma_{L}^{2}$ and $\mu_{L}$ are associated to variational distribution over $S$.

Recall in a regression task we have

that gives

Therefore, the likelihood function has form

Based on [4], we assume that the variational posterior of all parameters is a diagonal Gaussian distribution, i.e., $W|\mu_{w},\sigma_{w}^{2}\sim N(\mu_{w},\sigma_{w}^{2}I_{|W|})$, where $|W|$ denotes the number of weights and biases in the NN, and $\sigma_{w}=\log\left(1+\exp(\rho_{w})\right)>0$. Also let

such that

Further, assume $\varepsilon_{w}\sim\mathcal{N}(0,I_{|W|})$ and $\varepsilon_{L}\sim\mathcal{N}(0,1)$. Then, the posterior samples of $W$ and $S$ are given by

where $\odot$ denotes element-wise multiplication. Here, $\mu_{w}$ and $\rho_{w}$ are the variational parameters associated with $W$, and $\mu_{L}$ and $\rho_{L}$ are the variational parameters associated with $S$.

Algorithm 1 describes the VB algorithm for the regression network with variance uncertainty. We denote this method by VBNET-SVAR throughout the remaining of the paper.

In this section, we evaluate the results of our proposed model on two regression problems under the following scenarios

• •

  The normal prior, which corresponds to a normal distribution, assuming independency between the prior over parameters, is given by

  $$p(W)=\frac{1}{\sqrt{2\pi\sigma_{p}^{2}}}\exp\left(-\frac{1}{2\sigma_{p}^{2}}vec(W)^{\top}vec(W)\right),$$

  where $\sigma_{p}^{2}$ denotes the variance of the prior, and $\theta$ represents the vector of parameters.

• •

  The spike-and-slab prior [8], which corresponds to a mixture of normal distributions, assuming independency between the prior over parameters, is given by

  $$p(W)=\prod_{j}\left[z_{j}\mathcal{N}(W_{j};0,\sigma_{1}^{2})+(1-z_{j})\mathcal{N}(W_{j};0,\sigma_{2}^{2})\right],$$

  where, $W_{j}$ represents the $j$th (weight or bias) parameter in the model, $\sigma_{1}^{2}$ denotes the variance of the first mixture component, which is greater than $\sigma_{2}^{2}$, the variance of the second mixture component, and

  $$z_{j}\stackrel{{\scriptstyle\rm\tiny indep.}}{{\sim}}{\rm Ber}(\pi).$$

  The second variance $\sigma_{2}^{2}$ is assumed to be small ($\sigma_{2}^{2}\approx 0$), allowing the prior to concentrate samples around 0 with probability $1-\pi$.

We compare the performance of our proposed model with various models, including the model introduced in [4], and the frequentist NN. The results indicate that variance uncertainty can significantly enhance performance. It should be noted that in the nonlinear function estimation study, we have set the fixed variance for the model proposed by [4] to the train set MSE of the frequentist NN, while in the riboflavin data analysis, the fixed variance is set to the maximum of $0.2{\rm var}(y_{\rm train})$ and the train set MSE of the frequentist NN to prevent the overfit small error problem caused by the frequentist NN. These values are chosen by trial and error to obtain the best performance of the model of [4] (instead of a cross-validation approach).

In this paper, we considered variance uncertainty in neural networks, specialized for the regression tasks. Based on the MSPE and coverage probabilities, a simulation study of a simple function approximation problem and a real genetic data set analysis showed that the prediction performance of the BNN improved by considering variance uncertainty. This method applies when no information is available regarding the variance setting of the likelihood function, which is common in many real-world applications. This model can also generalize BNNs in regression problems with fixed variance. The code for this study is available online at github (https://github.com/mortamini/vbnet).