Posterior Inferred, Now What? Streamlining Prediction in Bayesian Deep Learning

Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-09-05 14:07:00

Abstract: The rising interest in Bayesian deep learning (BDL) has led to a plethora of methods for estimating the posterior distribution. However, efficient computation of inferences, such as predictions, has been largely overlooked with Monte Carlo integration remaining the standard. In this work we examine streamlining prediction in BDL through a single forward pass without sampling. For this we use local linearisation on activation functions and local Gaussian approximations at linear layers. Thus allowing us to analytically compute an approximation to the posterior predictive distribution. We showcase our approach for both MLP and transformer architectures and assess its performance on regression and classification tasks.

Main paper: 
 introduction: 
 Through the success of machine learning models in real-world applications, ensuring their reliability and robustness has become a key concern. In particular, in applications such as aided medical diagnosis [1] , autonomous driving [15] , or supporting scientific discovery [17] , providing reliable predictions, identifying failure modes, and identify how to reduce uncertainties of the system is vital.Uncertainty quantification is at the core of these topics with Bayesian deep learning (BDL, [20, 16] ) providing a promising paradigm for assessing uncertainties effectively and efficiently.The central goal in BDL is to make inferences w.r.t. the posterior distribution over the probabilistic model (the parameters or the function itself). For example, to compute the expected prediction, estimate model uncertainties, or use it within acquisition functions in active learning. For this, we need to first estimate the posterior distribution and secondly make inferences of interest based on the estimated posterior. While both of these steps typically involve intractable integration, only the first step has seen significant progress in recent years [2, 14, 3] . For the second step, in case of a Laplace approximation (LA, [10] ), globally linearising the model function around the maximum a posteriori (MAP) estimate to perform inferences [12, 8] has shown promise in providing good predictive uncertainty.However, for all other posterior approximation methods, sampling based approximations remain to be the default. Given the high dimensionality of neural networks, sophisticated sampling methods are usually computationally prohibited and vanilla Monte-Carlo sampling is typically employed.In this work, we tackle this problem by streamlining the prediction in BDL through local linearisation of activation functions and by utilising local Gaussian approximations at linear layers. Instead of a sample based approximation, which requires multiple re-evaluations of the network, we analytically approximate the posterior predictive distribution in a single forward pass through the network, making our methods well-suited for large-scale applications. Moreover, in contrast to global linearisation, our method is suitable for more complex inference tasks as the neural network function becomes locally linear with respect to the inputs. Empirically, we find that local linearisation and local Gaussian approximation of neural networks to provide accurate predictive uncertainties and predictions, while We evaluate our method on regression and classification tasks and find that our method result in good predictive performance. 
methods: 
 We denote the model parameters as θ and the training set as D. Given the inferred approximate Local Gaussian Approximations for Linear Layers Denote the weight and bias of the l th linear layer as W (l) ∈ R Dout × Din and b (l) ∈ R Dout respectively, and its input as a (l-1) ∈ R Din . Then the output h (l) is given as h (l) = W (l) a (l-1) + b (l) , where we use hposterior q(θ | D) = N (E [θ] , Σ θ ),(l)k to denote the k th element.We drop the superscript if it is clear from the context. Assumung that W , b and a are Gaussian distributed, we make the following two assumptions to obtain a tractable approximation on h: (i) we assume a i W ki is Gaussian, and (ii) we assume a i and W ki are uncorrelated.Under assumption (i), as the sum of Gaussian random variables (a i W ki ) is still Gaussian, h k will be Gaussian as well. Conseuqently, h will be jointly Gaussian distributed. The mean of h is given asE [h] = E [W ] E [a] + E [b]and the covariance between the k th and the l th hidden unit is computed as follows:EQUATIONwhereCov[a i W ki , a j W lj ] = E [a i ] E [a j ] Cov [W ki , W lj ] + E [W ki ] E [W lj ] Cov [a i , a j ] + Cov [a i , a j ] Cov [W ki , W lj ] .(2) Note that structure of the posterior covariance influences the computational cost of the approximation.Local Linearizations of Activation Functions Let g(•) denote a non-linear activation function computing a = g(h) for an input h. Given h ∼ N (E [h] , Σ h ), we use a first order Taylor expansion of g(•) at the input mean E [h] to obtain a tractable approximation of the distribution over a, i.e.,EQUATIONJ g | h=E[h] is the Jacobian of g(•) at h = E [h].As Gaussian distributions are closed under linear transformations, now a will also be Gaussian distributed, i.e.,EQUATIONW (l) × a (l-1) ≈ h (l) g(h (l) ) ≈ a (l)Local Linearization Eq. ((a) Multi-layer Perceptron (MLP) Note that the quality of the local linearisation will depend on the scale of the distribution over the input h. Combining local Gaussian approximations for linear layers and local linearisation for nonlinear activation functions results in a tractable approximation to the posterior predictive distribution. Fig. 2 illustrates our streamlined prediction for multi-layer perceptrons (MLP) and attention blocks in tranformers, for a detailed description on the approach for transformers see App. C.6. Note that the mean and covariance of the posterior predictive distribution can be computed in a single forward pass.H (l-1) t1 t2 t3 E H (l-1) W (l) Q E H (l-1) W (l) K H (l-1) W (l) V softmax( Q (l) K (l) ⊤ √ D )V (l) Q (l) K (l) V (l) (b) Attention BlockCovariance Structure Computing the full covariance of the posterior is usually infeasible due to high computational and memory cost. Diagonal approximation and Kronecker-factorization of the covariance/precision are two of the most common aprpoaches. For diagonal covariance, calculating the posterior predictive distribution is straightforward, see App. C.2 for details. In case of Kronecker factors, we developed a tailored block retrieval method for efficient propagation of uncertainties, see App. C.3 for details. Note that other covariance structures can exploited in a similar fashion.Computational Complexity We will briefly discuss the computational complexity of our method for the case of full covariance. Observe from Eqs. ( 1) and ( 2) that the computational cost to obtain(Cov[h k , h l ]) is O(D (l) in 2). Therefore, computing the output covariance at the l th linear layer will be in the order of O(D(l) out 2 D (l) in 2). For element-wise activation functions, the computational cost will be O(D (l) out). Hence, we obtain a total cost of O(L l=1 D (l) out 2 D (l) in 2 ) + D (l) out 2) for a network with L layers. By explointing the covariance structure, the total computational cost can be substantially reduced. 
experiments: 
 We evaluate our method on regression and classification tasks. We choose the Laplace approximation (LA, [3] ) and mean-field variational inference (MFVI, [18] ) to esimate the posterior. We compare predicitions based on sampling, global linearisation (GLM, [8] ), and our method. We use a paired t-test with p = 0.05 and bold results with significant statistical difference. For our method, we addionally fit a scale factor, multiplied to the predictive variance, by minimizing the negative log predictive density (NLPD) on the training set. This is necessary, as the predictive variance in case of deep and wide network with diagonal covariance structure can be large.Regression We experiment with multi-layer perceptron (MLP) for regression. See App. D.1 for experiment setup details and additional results. We use full covariance for LA. As shown in Table Table 1 , for MFVI our proposal (Ours) result in better performance than sampling on 8 data sets and matches the performance on the remaining 3 data sets. For LA, our approach obtains better performance than sampling on all data sets.Classification We train an MLP from scratch and fine-tune a pre-trained Vision transformer (ViT) base model [5] . See App. D.2 for experiment setup details and additional results. With LA, we use a Kronecker-factorized covariance for MLPs and a diagonal covariance for ViT models. As shown in Table 2 , we obtain better performance when compared with sampling and GLM. For ViT, fine-tuning on SVHN with MFVI failed, resulting in unreliable results. To test our method on out-of-distribution (OOD) data, we first evaluate the MNIST trained model on rotated MNIST as shown in Fig. 3a . We observe that with increasing roation degree, the increase in NLPD is less compared with other methods. In addition, we show OOD results on a FMNIST trained  
discussion & conclusion: 
 In this work, we proposed to streamline prediction in Bayesian deep learning by local linearisation and local Gaussian approximations. For this, we discussed the propgation in different neural network architecures and covariance structures. We showed through a series of experiments that our method obtains high predictive performance, obtain good predictive uncertainties, and can distinguish between in-distribution and OOD data. In future work, we aim to extend our approach to other network architectures, such as convolutional layers, and utilize our approach in more complex inference tasks. 
 
 Appendix: We first introduce notation in App. A and related work in App. B. Then, we introduce the derivation of our method in App. C. At last, we describe the experiment setup and additional experiment results in App. D.We use lowercase bold letter for vector, e.g., x, and uppercase bold letter for matrix, e.g., W . We use subscript to denote element of vector and matrix, e.g., x i (i th element) and W ki (k th row, i th column).We use W [k, :] to indicate the k th row of a matrix.Inferring Posterior in Bayesian Deep Learning There has been many methods developed which can be roughly grouped into three categories: (i) Laplace approximation based methods: Starting from [10] where simple post-hoc Laplace approximation (LA) has shown promising results on neural network, LA has gained increasing attention ever since. [4] has shown that treating a subnetwork Bayesian will also result in good predictive uncertainties. (ii): Variational inference (VI) based methods: [2] showed mean-field VI (MFVI) could improve generalisation in small-scale neural network and [18] showned MFVI is effective for large-scale neural networks as well. (iii): Others: Monte Carlo Dropout [6] aims to estimate predictive uncertainty by interpreting dropout in neural networks as a form of Bayesian approximation. Deep ensemble [11] combines the outputs of multiple independently trained models to capture predictive uncertainty. Stochastic Weight Averaging-Gaussian [14] , which extends Stochastic Weight Averaging [9] by capturing the posterior distribution of model weights using a Gaussian approximation.Making Prediction in Bayesian Deep Learning Little work has been done for this and the usual go-to solution is simple Monte Carlo Estimation. For Laplace approximation, [8] proposed a global liberalised model for better posterior prediction.We derive the approximate posterior predictive distribution form in this section. We make the following two assumptions to obtain tractable distribution on the output:• Assumption 1: We assume a i W ki is a Gaussian distribution.• Assumption 2: We assume a i and W ki are uncorrelated.From assumption 1, given W , a, and b are all Gaussian, each h k will be a Gaussian distribution. As a result, h will be a Gaussian distribution as well.We now derive its mean and covariance. We first derive the mean for each h k . As a i and W ki are uncorrelated, we haveEQUATIONEQUATIONEQUATIONWe now derive the covariance. DefineEQUATIONwhereEQUATIONEQUATIONTo derive the form of Cov[a i W ki , a j W lj ], we use assumption 2:EQUATIONNow the only term left is Cov[h ′ k , b l ], which can be written asEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONEQUATIONPutting it together, we haveEQUATIONwhereEQUATION)Note that the first term in Eq. ( 24) could be rewrite into the form of matrix multiplication which results in an efficient implementation: EQUATION+    E [W k1 ] E [W l1 ] . . . E [W k1 ] E [W l Din ] . . . . . . . . . E [W k Din ] E [W l1 ] . . . E [W k Din ] E [W l Din ]    ⊙    Cov[a 1 ,EQUATIONWhen the posterior has diagonal covariance, the mean E [h k ] will still be the same.For covariance, note that as now the posterior is diagonal, when k ̸ = l, we haveEQUATIONFor k = l, we haveEQUATIONIn Kronecker approximation, the Hessian is represented in Kronecker product form: Table 5 gives the classification data sets information and the neural network structure we used. We use ReLU activation for MLP. For ViT, we make the MLP block in the last two transformer block and the classification head Bayesian, and treat the rest of the weight deterministically. In Table 6 we report the test accuracy, on SVHN the fine-tuning of MFVI failed and hence the bad performance.For the rest, Ours results in matching or better performance compared with sampling and GLM, indicating the effectiveness of our method.  



Reviews: review: title: title: An okay submission, but not fully convinced about the need for the method or its performance
main_review: main_review: ## Main review

In this paper, the authors propose to propagate the uncertainty that comes from the posterior over the weights of a Bayesian neural network by applying an approximation to the weight-activation products and a local linear approximations to the activation functions.

This approach is interesting, but it is also not fully new (see https://arxiv.org/pdf/1810.03958). The experiments seem reasonably thorough, but the report of error bars (and marking of "winning") methods is unclear. Generally, I am not convinced that the proposed method offers a clear performance advantage over, say GLM, or why one would need this approach in the first place, compared to simply using GLM. A clearer framing of what the proposed approach offers that existing linearisation methods do not is important here.

## Suggested minor adjustments

L2-3: “efficient computation of inferences”: Phrasing here is a little off. Predictions are not typically viewed as something that is inferred.

L7: “layers. Thus allowing…”: probably a comma instead of a full stop.

L14: “identifying failure modes, and identify”: identify repeated.

L19-20: “ function itself). For example…”: writing here reads a bit off as well. Maybe comma instead of full stop?

L20-21: “secondly make inferences of interest based on the estimated posterior”: again, as per one of the previous comments, we typically make inferences about the posterior itself or, in general, latent variables, not downstream quantities (predictions) as seems to be implied here.

L25-26: “remain to be the default”: ungrammatical. Maybe “remain the default”?

L27: “prohibited”: ungrammatical. Maybe “prohibitive”? Also, your wider argument here that “sophisticated sampling methods are usually computationally prohibitive” is not very convincing, and virtually none uses vanilla Monte-Carlo for integrating out the posterior distribution.

Figure 1: The authors claim their method’s predictive uncertainty is somehow “better” but do not support this with an argument. I don’t think this obvious from the figure.

L43: “our method result”: typo, should be “results.”

L46: the notation E[theta] is unclear. What distribution are you taking the expectation with? Why not just use a parameter, e.g. mu, here?

L56: the “assumption” that “aW” is Gaussian isn’t an assumption, it’s an approximation. If a and W are Gaussian, their product is *not* Gaussian. I would rephrase this to say, that “we assume W, b, and a are Gaussian, and that a and W are uncorrelated. While aW is non-Gaussian, we will later approximate its distribution by a Gaussian distribution.” Also, the assumption that a and W are independent is not really an assumption: a and W really are independent (for most standard feedforward architectures), since the activations that are input to a layer do not depend on the weights of the layer.

L58: “Conseuqently”: typo, “Consequently”

Figure 3 (a): legend hides half of the plot, consider rearranging.

Figure 3 (b): histograms overlap each other and make it hard to read the plot. Also the histogram bars are unequal sizes.

Tables 1 and 2: good to see error bars, but what do they mean? Are these 95% confidence intervals on the NLPD? In most cases, your error bars overlap with the error bars of the sampling method, in which case you should probably boldface both methods.

L113: “propgation”: typo, “propagation.”

L115: “obtain”: typo, “obtrains” (also repeating obtain).
title: title: The overall presentation of the paper is good with quite broad experiment results. The paper's contribution is valuable, especially since the topic is still underexplored.
main_review: main_review: The paper proposes a new technique to approximate the predictive distribution in Bayesian Deep Learning (BDL) given an approximate posterior. The introduced technique is faster compared to the standard Monte Carlo (MC) sample-based method. The core idea is to make a Gaussian approximation for the neural network's linear layer and perform local linearizations on the activation function.

The overall presentation of the paper is good with quite broad experiment results. The paper's contribution is valuable, especially since the topic is still underexplored.

some comments:

1. Any comments on why only the last two transformer blocks and the classification head are considered here? For example, what are the trade-offs regarding speed and performance in this context? it might be good to mention this in the paper.
>we make the MLP block in the last two transformer block 282 and the classification head Bayesian, and treat the rest of the weight deterministically.
2. Currently, two types of neural network architectures are considered, both of which are widely used in various applications today. What are the possibilities of using other architectures for this? In other words, how general is the proposed method?
3. In Figure 3, IVON is mentioned, but I couldn’t find any reference to it in the main text. Is the IVON mentioned from Shen et al. (https://arxiv.org/abs/2402.17641)? I think this should be fixed in the camera-ready version.