Streamlining Prediction in Bayesian Deep Learning
Authors: Rui Li
Marcus Klasson
Arno Solin
Martin Trapp
Venue: ICLR 2025 Poster (2025)
Keywords: Bayesian deep learning
uncertainty quantification
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We streamline prediction in Bayesian deep learning through a single forward pass without the need of sampling.

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 of activation functions and local Gaussian approximations at linear layers. Thus allowing us to analytically compute an approximation of the posterior predictive distribution. We showcase our approach for both MLP and transformers, such as ViT and GPT-2, and assess its performance on regression and classification tasks.

Open-source library: https://github.com/AaltoML/SUQ.

Published as a conference paper at ICLR 2025
STREAMLINING PREDICTION IN BAYESIAN DEEP LEARNING
Rui Li
Marcus Klasson Arno Solin Martin Trapp
Department of Computer Science, Aalto University, Finland (fi rssnamme, lastname) ) aa 1to.
ABSTRACT
The rising interest in Bayesian deep learning (BDL) has led to plethora of meth- ods for estimating the posterior distribution. However, efficient computation inferences, such predictions, has been largely overlooked with Monte Carlo inte- gration remaining the standard. In this work we examine streamlining prediction BDL through single forward pass without sampling. For this, we use local lineari- sation of activation functions and local Gaussian approximations at linear layers. Thus allowing us analytically compute approximation the posterior predic- tive distribution. We showcase our approach for both MLP and transformers, such as ViT and GPT-2, and assess its performance on regression and classification tasks.
Open-source library: httpp::/gggteet com wow/AaltoMI//SO
INTRODUCTION
Recent progress and adoption of deep learning models have led to sharp increase in interest of improving their reliability and robustness. applications such as aided medical diagnosis (Begoli al., 2019), autonomous driving (Michelmore et al., 2020), supporting scientific discovery (Psaros et al., 2023), providing reliable and robust predictions as well as identifying failure modes is vital. A principled approach to address these challenges is the use of Bayesian deep learning (BDL, Wilson & Izmailov, 2020; Papamarkou et al., 2024) which promises a plug & play framework for uncertainty quantification. However, while plugging the Bayesian approach into deep learning is relatively straightforward (Blundell et al., 2015; Gal & Ghahramani, 2016; Wu et al., 2019), the play part is typically severely hampered by computational and practical challenges (Wenzel al., 2020; Foong et al., 2020; Gelberg et al., 2024; Coker et al., 2022; Kristiadi et al., 2023).
The key challenges associated with BDL roughly be divided into three parts: (i) defining meaningful prior, (ii) estimating the posterior distribution, and (iil) performing inferences of interest, e.g. making predictions for unseen data, detecting out-of-distribtioon settings, or analysing model sensitivities. While constructing meaningful prior is an important research direction (Nalisnick, 2018; Meronen et al., 2021; Fortuin et al., 2021; Tran al., 2022), has been argued that the
Practical Outlier Detection
Input Sensitivity Analysis
f(x)
Domain
Out Domain
C
MAP Model
Entropy
Figure 1: Our streamlined approach allows for practical outlier detection and sensitivity analysis. Locally lincarising the network function with local Gaussian approximations enables many relevant inference tasks to be solved analytically, helping render BDL practical tool for downstream tasks.

---

Published as a conference paper at ICLR 2025
differentiating aspect of Bayesian deep learning is marginalisation (Wilson & Izmailov, 2020: Wilson, 2020) rather than the prior itself. Hence, estimating the posterior distribution has seen significant progress in recent years (Blundell et al., 2015; Maddox al., 2019; Daxberger et al., 2021a) with particular focus on post-hoc approximations (Kristiadi et al., 2020; Daxherger et al., 2021b). However, while these approaches have shown promise in making BDL. useful for real-world applications, they tackle only part of the computational and practical challenges associated with BDL.
In this work, we focus on streamlining prediction in BDL for downstream tasks by providing straightforward and effective method to compute inferences of interest, cf., Fig. 1. For this, we make the neural network locally linear with respect to the inputs. Thus, inferences, such as computing predictions, admit a closed-fore solution and can be estimated efficiently. In particular, we propose using local linearisation of non-linear activation functions at every layer of the network and local Gaussian approximations at linear layers. Empirically, we find that local linearisation combined with Gaussian approximation of Bayesian neural networks provides accurate predictions, with useful predictive uncertainties, while being conceptually simple. Moreover, complex inference tasks w.r.t. the inputs, such as analysing model sensitivities to input perturbations, can be computed efficiently. Thus allowing us truly account for all sources of uncertainties.
Contributions: (i) We propose layer-wise local linearisation and local Gaussian approximations of neural networks to streamline BDL for downstream tasks (Sec. 3). (ii) We discuss how to handle different covariance structures and architecture choices (Sec. 3.2 & Sec. 3.3). (iti) Finally, we present an empirical assessment of our approach on regression and classification tasks, and showcase its utility for uncertainty quantification, out-of-domain detection, and sensitivity analysis (Sec. 4).
RELATED WORK
To estimate the posterior in BDI., variational inference (VI, Blei et al., 2017; Zhang et al., 2018) utilises a variational approximation to the true posterior distribution and minimises a divergence measure between both distributions. A typical choice for the variational family Gaussian distribution, chosen for computational reasons. Early works on mean-field VI (MFVI) and related approaches require modifications of the model structure (Blundell et al., 2015) to perform a reparametrisation of the variational distribution. Recent work by Shen et al. (2024) developed an optimiser to ease the use of MFVI, and has shown good performance on large-scale models such as ResNets (He et al., 2016) and GPT-2 (Radford et al., 2019). However, VI-based methods typically require Monte Carlo estimation to perform inferences, which can be problematic in practice due to additional computational overhead.
factorised
A recent trend in BDL are post-hoc methods, such as the Laplace approximation (LA, MacKay, 1992a), which can be applied directly on the trained model without modification (Kristiadi et al., 2020; Daxberger et al., 2021a). Daxberger et al. (2021b) extended the applicability of LAs by showing that treating a subset of parameters Bayesian can still give good predictive uncertainties. Moreover, Immer et al. (2021b) proposed the linearised LA by performing a global linearisation, which is principled under the Generalised Gauss approximation to the Hessian, and has shown promise in providing useful predictive uncertainties. Recent works applied post-hoc methods in various applications. such as large language models (Yang et al., 2024; Kampen et al., 2024), vision-language models (Baumann et al., 2024), dynamic neural networks (Meronen et al., 2024), and sequential learning (Scannell et al., 2024).
In addition, various tailored ensemble- based methods for BDL have been proposed, such as Monte Carlo dropout (Gal & Ghahramani, 2016), deep ensembles (Lakshminarayanai) et al., 2017), and stochastic weight averaging- Gaaiisinn (Maddox et al., 2019). While some works on deep ensembles enable estimating the predictive distribution in single forward pass (Eschenhagen et al., 2021; Havasi et al., 2021), most methods typically require multiple forward passes to estimate the predictive distribution and do not explicate an approximation to the posterior distribution.
More recently, there has been trend in exploring deterministic computations in BDL to avoid the need for sampling (Goulet et al., 2021; Giordano et al., 2024: Burroni et al., 2024). In particular, Wu et al. (2019) derived an analytical training objective for VI VI using moteen-matching each layer of the network. However, the solutions to the moment-matching have to be derived manually for each type of activation function, making it impractical in practice. More recently, Goulet et al.
2

---

Published as a conference paper at ICLR 2025
Local Gaussion Local Lincarisation Approxination
Qi wettmor 220
(V10
[[
K
(1-1)
9(h
(a) Multi- Perceptron (MLP)
(b) Attention Layer
Figure 2: Illustration of our approach for different network architectures. In MLPs, we can directly apply local Gaussian approximations and local linearisation of each layer. The distribution over activations then propagated to the next layer. attention layers, we treat the query Q and key K deterministically and only treat the value V as random quantity, resulting in a straightforward propagation path. The resulting distribution then propagated to the subsequent MLP layer.
(2021) proposed local linearisation of the network to perform message passing on the network under mean-field assumption. Moreover, Petersen al. (2024) used a local lincarisation of the network to propagate alcatoric uncertainties over the input through a deterministic network. In addition, Dhawan et al. (2023) investigated local linearisations of activation functions to estimate the function space distance of two neural networks, for example, relevant in continual learning settings. In contrast, our work disentangles the approximation of the posterior distribution and the computation of inferences w.r.t. the posterior distribution. Hence, it provides a streamlined framework to propagate all forms of uncertainties through Bayesian neural networks.
3 METHOD
In Bayesian deep learning (BDL), predicting the output y (e.g., class label, regression value) for an input æ E X' is performed by marginalising out the model parameters 0 of the neural network fe(-) instead of trusting a single point estimate, i.e.,
1/ ply pl de,
p x. D)
(1)
where D (I (n)) = denotes the training data and the posterior distribution p(0 (D)
2(0.D)
p(D)
is given by Bayes' rule. However, for most neural networks, integrating over the high-dimensional parameter space is intractable, necessitating the use of approximations to compute the posterior distribution p(0) D) and the posterior predictive distribution p(v x. D).
Recently, much progress has been made in efficiently approximating the posterior distribution for BDL, including scaling mean-field variational inference (Shen et al., 2024) to large-scale models and performing post-hoc estimation using the Laplace approximations (Daxberger al., 2021a). A common thread is using tractable distribution to approximate the posterior distribution g(0) z | commonly chosen as Gaussian distribution. Consequently, the posterior predictive distribution is typically approximated using Monte Carlo integration, i.e., by sampling from to estimate the integral in Eq. (I), with the exception of the linearised Laplace approximation (Immer et al., 2021b). However, while using Gaussian approximation facilitates efficient computation of the approximate posterior distribution, sampling from the high- sigh-dimensional Gaussian approximation can be challenging (Vono et al., 2022) and result in high computational overhead.
We will now shift our focus on estimating integrals of the form of Eq. (1) and assume that an approximation to the posterior distribution g()) is given. Further, we will assume that g is in the family of stable distributions. Note that a linear combination two independent random variables following stable distribution has the same distribution as the distribution of the individual random stable distribution. Marginalisation
variables. The Gaussian distribution is a typical example tasks such as in Eq. (1) appear in many scenarios, e.g., active learning (MacKay, 1992a: Gal et al., 2017; Smith et al., 2023), model selection (Immer et al., 2021a; MacKay, 1996). or outlier detection (Wilson & Izmailov, 2020), and pose reappearing challenge in downstream applications of BDI..
3

---

Published as a conference paper at ICLR 2025
3.1 STREAMLINING COMPUTATIONS WITH LOCAL APPROXIMATIONS
Let the weights and biases of the m linear layer of the network / be denoted as W(m) e Raax Du. and b(m) E IR IR,,,, respectively. Then the pre-activation (m) is given as h (m)) (m)), a (m-1) + b(m), where a (m--1) E Da is the activation of the previous layer. In case
m then (0)) corresponds to the input æ. We further denote the Ethh element of h (m) as Da k = DD i= Wh d (m-1)) + [(m) and drop the superscript if it clear from the context.
Given an upproximate posterior distribution (()) with 0 = { (W(m), b
J==1, we aim to compute the probability distribution of the activation a (m) of each layer m. For this, we need to estimate the distribution of the pre-activation h (m)) and then compute an approximation to the activation a (m) after the application of non-linear activation function g(-).
Approximating the pre-activation distribution In case the activation !!"-1) is deterministically given, Le., for the input layer, we can compute the distribution over pre-activations analytically as a consequence of the stability of stable distributions under linear transformations (Petersen et al., 2024). However, for hidden layers, the distribution over pre-activations is generally not of the same family the posterior distribution (Wolinski & Arbel, 2022). Nevertheless, we will apply a local Gaussian approximation to the pre-activation at every hidden layer. Specifically, we make the assumption:
Assumption 3.1. Assume that the activations of the previous layer dm-1)) and parameters of the th layer are independent.
Then followed by Gaussian approximation of a pre-activation h (m)) is given as:
(m-11) w-(m) for each i and each ks the mean of the
(m)
E (m)
E a (m-1))
1)]
E (m))
(2)
and the covariance between the 1.th and the jth hidden unit is computed as:
Cov
h
(m²
A, (m))
Σ Cov
[ (a.m--11 (m²) ks
:a
(m
(im-1/w(mm Wr(mm) + Cov [B(m), ((fm))))))
(m=-
(Cov [ (mm) + Cov [w/mm), 6.6% b6ml))
a;
(3)
where
Cov
a
(w(())
.
(m-
1) 1)
(m- 1)
E
a
(m- 1)
Cov [w/m), wr(m)
ki
(m())
(m)
(m)
(m-1) (m-1)
+
E
W
1j
Cov
a; aa
+ Cov
a
(m
(m
Cov
[w(m), wm(m)
(4)
A detailed derivation alongside empirical evaluation of the approximation quality can be found in Apps. A.I and A.2. Depending on the structure of the covariance matrix, we can further simplify the computation of the covariance matrix, which we will discuss in Sec. 3.3.
Approximating the activation distribution Let g(-) denote a non lon-linear activation function computing g(h) for pre-activation h. Inspired by the application of local lincarisation in Bayesian filtering (e.g., Särkkä & Svensson, 2023), we use a first-order Taylor expansion of g(-) at the mean of the pre-activation E [h]. Specifically,
approximate ghh) using
(5) where J Joln- is the Jacobian of g(-) at h E [h). Then, as stable distributions are closed under linear transformations, the distribution of can be computed analytically and is given as follows in case of a Gaussian distribution, i.e.,
g(h) = g (E [h]) + J, ln B. (h E [h])
(E (1). J,yh T/a En l=jaa)
Note that the quality of the local linearisation will depend on the scale of the distribution over the input h. For ReLU activation functions, Petersen et al. (2024) have shown that local linearisation
4

---

Published as a conference paper at ICLR 2025
provides the optimal Gaussian approximation of univariate Gaussian distribution in total variation. For classification tasks, we employ a probit approximation MacKay (1992b); Kristiadi et al. (2020).
One way to understand the resulting approximation is as piecewise linear function (or multilinear function). Globally, the function will still be non-linear, but locally it will behave linearly. In contrast to the original model, which composes piecewise linear functions in the case of a ReLU network, our approximation composes linear functions locally. We obtain a piecewise linear function due to the local composition, which allows us capture the non-linear nature of the model.
Intuition
3.2 ARCHITECTURE CHOICES
By combining local Gaussian approximations for linear layers and local linearisation for non-linear activation functions, we can analytically compute the distribution over activations at each layer in single forward pass. In the case of a multi-layer perceptron (MLP) and common architecture choices, the described approach can be directly applied to each layer the network. However, further considerations needed to streamline the computation path for more complex architectures such as attention. Fig. 2 illustrates the computation path for MLPs and attention layers.
Attention layers Each block in transformer (Vaswani et al., 2017) constitutes: multi-head attention, an MLP, layer normalisation, and a residual connection. For the MLP part, the propagation the same as previously described. Further, layer normalisation is linear transformations, and the resulting distribution can be obtained analytically. For residual connections by assuming independence, we could also obtain the resulting distribution analytically. Treating the multi-head attention block is more involved as the softmax activation function 'squashes' the distribution of the pre-activations. We describe our method below, and further details are App. A.6.
Given an input H E RT'D, where T' is the number of tokens in the input sequence and D is the dimension of each token, denote the query, key and value matrices as WQ E DxD WK E DDDD Wv IR Dx respectively. Further, we denote the key, query and value in an attention block as Q = HWQ. K = HWK, and V HWv. Then the output of attention layer is given as follows Attention| H) = Softmax (QKT/ /VD) V. For computational reasons, we will assume the input distribution to the multi-head attention block has diagonal covariance structure. 'Pushing' random vectors over softmax activation may require further approximations and will not result in an output with a distribution close to a Gaussian distribution. Hence, we treat the query and key matrices as deterministically given. A possible remedy is to leverage an approximation to the softmax function such as Lu et al. (2021). Consequently, the attention scores are given as:
IE [ E W [B]H] [ W V,
Attention(H) Softmax
VD
where V follows stable distributions. Due linearity, the resulting distribution again be obtained analytically.
3.3 COVARIANCE STRUCTURE
Computing the full covariance the posterior is usually infeasible due to high computational and memory cost. We describe our methods for the two most common covariance approximations and will briefly discuss the computational cost in the case of a full and diagonal covariance structure.
Full covariance When the posterior has full covariance, for the mh linear layer the computational complexity for computing Cov(hy, hu] is ((D((m)) Consequently, computing the covariance of the activations for the mh layer adds to O(D (m) [TD (m))). Computing the local linearisation for element- wise activation functions results in complexity of ((D!!i)))) Hence, we obtain total cost of O .M ( + {D<mr) (2) ) for network with M layers. As the computational cost is directly linked to the number of parameters and their correlation structure, a natural way to reduce the computational cost is to either exploit the structure in the covariance matrix consider only subset of parameters, in the spirit of subnetwork Laplace (Daxberger et al., 2021b). We will focus on exploiting the structure of the covariance using subset of parameters trivially extends from discussion.
out
5

---

Published as a conference paper at ICLR 2025
Wi1, WasT WT., W23
Wn1, Wsq
WT1. Wsi
Wn. W WoPTT
Wi
W33
Wiz, Wix Wia. W12
W12, Wis
W1a, Wsi
W12. Ws2 W1a.
Wa3
Cor[W]
Wis, Wopy W1a
W13, W137
11,, Wa1
W13, Way 1 Va3, Waa
W21, Wi1 W21, W12
W21, W13 W21, Ws1
W21, Wa2 W21, W33
W22, Way 222, 122, W12
W22, Wiyg 22, Waa
W22, W Wapx W2a, Waa
W2s, Wi1 W2s, W12
W2s, W13 Was, Wsi
W2s, Ws2 Was, Was
retrieve the highlighted submatrix Cov[W[[, [2, of covariance for W E R2x3, we identify the Kronecker blocks that contain the covariance of interest (II, III, V, and VI), explicate those blocks in memory, and then retrieve the relevant submatrix.
Figure 3:
Diagonal approximation In case the correlations between model parameters are ignored, as in mean-field variational inference, the computation of the pre-activation covariance reduces to:
(m) (m( (m)] = Σ E W
[w][m]] Cov [am-= (m-1) (m--)) 1)
Cov
(8)
DI
and variance of the hhhh pre-activation is given as: Varlh (m]
ΣB EB(omml (m) Var (w((m)) + Var (mm))) Var (amm-1)) (E [m)) ?
Var (
ISiSD.
Hence, assuming a diagonal covariance structure can help in reducing the computational burden. If only the variance the layer output of interest, the computational cost can be further reduced and adds to a total of O M (mm) (m)) run time analysis indicating little to no overhead given App. B.7.
(mm) D(T)) + D(m)). Further details are given in App. A.3. An empirical
(KFAC) Another common choice for approximating the posterior co- variance the use -ronecker-aacttoiistioon (KFAC) (Martens & Grosse, 2015), popularised in the context of Laplace approximations (Ritter et al., 2018). In this Σ is given by a Kronecker product of two factors, i.e., Σ (A & B + 1²1) where @ denotes the Kronecker product and I is prior precision. Note that in case of non-zero prior precision, the covariance cannot be expressed in the form of a Kronecker matrix multiplication. Denote the eigenvectors and eigenvalues of as UA and A. respectively, we approximate the Σ as follows:
Pronecker-factorisee
the posterior covariance
& B+ (UU (UBAAUU)) + AZT) ((U, (A + AL)) ((A n AL,) @U))T) (10)
(Eigen Decomposition)
To compute the covariance of the pre-activations, we need to retrieve the covariance between the weights of the 4h unit and the weights of the tth unit, which corresponds the h and 7th row in W, i.e, Cov [ov[W[[: :k. :[. In case of KFAC Laplace approximations, accessing Cov[W[[: [k. :. W [I] cannot be done directly. Therefore, we developed block retrieval method to retrieve Cov W [k. W[.:]] without explicating the full covariance matrix in memory.
The key idea is first identify the Kronecker blocks that contain the covariance of interest and then retrieve the submatrix by reconstructing only the relevant blocks. Fig. illustrates the idea in a toy example for weight matrix W R2x3 and submatrix of interest Cov [ [1, W :I]. Our method only reconstructs blocks containing the sub-covariance of interest (II, III. V, and VI) and then retrieves the relevant submatrix. Further details are given in App. A.4.
EXPERIMENTS
We demonstrate practical applicability of our approach classification/eeressson tasks (Sec. 4.1), large-scale classification results with VIT/GPT models (Sec. 4.2), and sensitivity estimation (Sec. 4.3). Additional experiments and additional experimental results can be found in App. B.
Data We use selection of data sets from the repossirryy (Kelly et al., 2023) for the regression experiments. For classification, experiment MNIST (LeCun et al., 1998), FMNIST (Xiao et al.,

---

Published as a conference paper at ICLR 2025
2017), well as the 11-class data sets OrganCMNIST and OrganSMNIIST from MedMNIST (Yang et al.. 2023). To assess our method on higher- fisher-dimensional settings, we experiment with CIFAR- 10 and CIFAR-100 (Krizhevsky & Hinton, 2009). DTD (Cimpoi et al., 2014), RESISC (Cheng et al., 2017) and a subsampled version of ImageNet-R (Hendrycks et al., 2021) with 100 classes to reduce the memory overhead for the LA. For the GPT model, used the BOOLQ. WIC, and MRPC tasks from GLUE (Wang et al., 2019b) and SuperGLUE (Wang et al., 2019a) benchmarks.
Posterior approximations We adopt the Laplace approximation (LA, MacKay (1996) and mean- field variational inference (MFVI, Blei et al., 2017) for approximating the posterior distribution of the network parameters. For the LA, we estimate the full covariance for the regression experiments, while we use diagonal or KFAC approximations for the covariance where applicable in the classification experiments. We compare our method using local Gaussian approximation and local linearisation against Monte Carlo (MC) sampling and global linearised model (GIM, Immer et al., 2021b). For MFVI, we adopt the IVON optimiser (Shen al., 2024) to obtain the posterior approximation with diagonal covariance structure by default, which has been shown to be effective and scalable to large-scale classification tasks. Here, we compare our method against MC sampling from the posterior to make predictions as done in Shen al. (2024). For the MFVI and LA sampling baselines, we used 1, 000 MC samples in the regression and classification experiments in Sec. 4.1, and 50 MC samples for the ViT and GPT-2 in Sec. 4.2. For method, fit an additional scaling factor the predictive variance by minimising the NLPD on a validation set, similar to the pseudo-count used in Ritter et al. (2018).
Network architectures We experiment with one or two-layer multi-layer perceptron (MLP) on the UCI regression data sets with details given in App. B.1. For MNIST, FMNIST, OrganCMNIST and OrganSMNIST, we use an MLP with layers containing 784 128 64 - C' neurons, where C' is the number of classes. For CIFAR- 10/100, DTD, RESISC and ImageNet-R, we fine-tune Vision Transformer (ViT) (Dosovitskiy et al., 2021) base model pre-trate pre-trained on ImageNet-- lk (Deng et al., 2009). For the GPT model, we use the pre-trained GPT-2 base model from Hugging Face Transformers (Wolf et al., 2019) and fine-tune it on the respective tasks.
Evaluation metrics For the regression experiments, we measure the negative log predictive density (NLPD) and root-mean-square error (RMSE) for each method. In the classification experiments, we use accuracy (ACC), NLPD, and expected calibration error (ECE) to compare the methods. We use paired t-test with p 0.05 to bold results with significant statistical differences when reporting the results. For assessing out-of tit--ffdiissiit distiriiibttionn (OOD) robustness, we use a Gaussian kernel density estimator with variance of 0.25 on the histogram of the predictive entropy evaluated on the test set.
4.1 DOES OUR METHOD PROVIDE USEFUL UNCERTAINTY ESTIMATES?
Regression We experiment on a selection of data sets from the UCI repository and run a 5-fold cross validation to report results for each data set. We use either MFVI or LA to obtain the posterior approximation and separately compare our method against their corresponding prediction approaches. Table I shows our method achieves better NLPD in general than the predictions with sampling for both MFVI and LA. Moreover, our method performs on par with the GLM, even though our method results in locally linearised network w.r.t. the inputs. Similar conclusions are made inspecting Table 8.
Table I: Negative log predictive density 4 on UCI regression data sets. Ours results in better or matching performance compared with sampling and GLM, indicating the effectiveness of our method.
MFVI (Diagonal Covariance)
Laplace Approximation (Full Coveriance)
(n, d)
Sampling
Ours
Sampling
GLM
Ours
SERVO
(167, 4)
1 20.000
1. 133-=008
3.795:2.110
1.017=0.172
1.113=0.007
5)
1.346::0.290
1.369-0.40
2.221::0.10
1.4744:0.64
AM
(398,
1.004
0.807=0.06
1.81 2::0.005
0.492=0:278
0.478
REV
(414,
1.076 029
0.92520.0m
1.93220.045
0.859
1.833
FF
(517.
2.160:3.00
2.333-3.61
2.086:20.292
1.584=0.550
1. 1.596
ITT
(1020, 33)
0.937 047
0.841 000
1.68 1:0.09
0.82520.005
0. 0.75620.116
CCS
(1030.85 8)
0.939 oc8
0.828=0.m 108
1.612+0.045
0.319:0.109
0. 0.233110.161
ASN
(1503, 5)
0.962=0.005 054
0.899=0006s
1.7888=0.045
0.4222+0100
0. 0.396
127)
0.973:00 062
0.920=0.118
1.84880.005
1.28120.060
2.662::1.00
PT
(5875, 19)
0.976 one
0.940=0.071
0.984
0.576200181
0.65 an
CCPP
(9568, 4)
0.365 040
0.352=0.04
1.3/15:0.055
000622-0182
-00622-0.200
Bold Count
3/11
11/11
0/11
7/11
8/11
7

---

Published as a conference paper at ICLR 2025
Table Performance metrics on the MNIST-like data sets for each method with the standard error for ACC and NLPD. Our method achieves better NLPD and ECE than the baselines.
Metries
Methods
MNIST
FMNIST
ORGANCMNIST
ORGANSMNIST
LA Sampling
0.97220.002
0--08
0.734:10.008
0.591:00011
LA GLM
0.975:20.00
0.882:0.003
0.824=0007
0.63440.000
ACC
Ours
0.975+000002
0.826=0.008
0.630=0.000
MFVI Sampling
0.974++0.00
0.620s0.010
0.167700010
MFVI Ours
0.97449.002
0.842
0.630=0.010
0,40420.010
LA Sampling
0.210
0.556-0.008
1.135:20.017
1.61420.021
LA GLM
0.089
015
0.87614.44
10.07
NLPD 4
LA Ours
0.089
0.397=0.010
0.710=0.02
1. 1.365=0.000 025
MFVI Sampling
0. 0.179=0.004
2. 2.010=0.o61
4.775+0.140
6.095=0.141
MFVI
LA Sampling
0.122
0.151
0.292
0,261
LA GLM
0.009
0.078
0.038
0.1
ECE
Ours
0.001
0.012
0.122
0.151
MFVI Sampling
0.018
0.144
0.251
0,385
MFVI Ours
0.003
0.013
0.145
0.067
Table 3: Performance comparison between our method and Moment-Matching (MM) on the MNIST- like classification tasks. We report the ACC and NLPD with standard errors and the ECE, Our method outperforms MM despite being simpler and more applicable to various distributions.
Metries
Methods
MNIST
FMNIST
ORGANCMNIST
ORGANSMNIST
MM
0z
0.88210.0
0.77
0.60020.01
ACC1
Ours
0.975
0.881 t0.005
0.816
0.614:0010
MM
U. 0.30000000
U. 004
D. 0.8
1.401=0.00 =0.0%
NLPD
Ours
s0.00
0.435
0.733
1.282+0.02c
MM
0.014
0.051
0.023
U.110
ECE
0.006
0.127
0.081
Here, we assess our method on MNIST-like classification tasks. For the LA, we use KFAC approximation of the covariance reduce the memory overhead. In Table 2, we report the ACC and NLPD with their standard eтrоrs and the ECE for each method. Our method achieves similar ACC with the baselines, while outperforming them on the NLPD and ECE metrics. In App. B.2, we assess our method on robustness to OOD data. We evaluate an MLP trained on MNIST on rotated versions of the test set. Our method consistently reduces overconfidence on OOD data, cf., Fig. 8.
Classification
Ours vs. moment-matching To verify the viability of local linearisation, we compare our method moment-matching
against moment- matching (MM) used in Wu (2019). We apply MM instead of local linearisation to our setting, assuming diagonal posterior approximation from LA. In Table 3, show the results for our method against MM on MNIST-like classification tasks. Our approach outperforms MM across the data sets and the metrics, except for the ECE on OrganCMNIST.. Compared to MM, our method is applicable to any differentiable activation function and any type stable distribution (Petersen al., 2024), while MM requires tailored derivations for each case and, hence, less plug-and-plan
4.2 IS OUR METHOD SCALABLE?
We demonstrate that our method is applicable to large-scale networks by experimenting with pre- trained ViT and GPT-2 models. In particular, we experiment with applying our method on either the attention layer or the MLP after the attention layer in the last two (VIT) four (GPT) transformer blocks. For each target data set, we fine-tune tune the layers we obtain the posterior approximation for. We assume a diagonal posterior to reduce the memory overhead. Table 4 shows the results when fine-tuning ViT models and obtaining the posterior approximation from the attention layers. We observe that our method achieves better or on par NLPD and ECE compared to the baselines for both LA and MFVI across all data sets while maintaining similar ACC as the basclines.
In Table we show the results for a GPT-2 model with LA on the MLP layers. We observe that our method systematically outperforms sampling, while achieving similar performance GLM in some cases. Thus indicating that our method applicable to different application domains. We present additional results for ViT models in App. B.
We also assess the robustness to out-of- distrribbttton (OOD) data for our method and the baselines. In particular, take the ViT network fine-tuned on CIFAR-10 and evaluate its predictive entropy on
8

---

Published as a conference paper at ICLR 2025
Table Performance metrics using ViT with posterior approximation on the attention layers with the standard error for ACC and NLPD. Our method achieves better NLPD and ECE in general and achieves similar ACC compared to the baselines.
Metrics
Methods
CIFAR-10
CIFAR-100
DTD
RESISC
IMAGENET-R
LA Sampling
0.971 20002
0.88210.005
0.715-10.00
0.89220.001
20.012
LA GLM
0.976
0.879=000
0.71810.000
0.091
0.73920.112
ACC1 T
LA Ours
0.976
0.880+0.00
0.719
0.89220.001
0.73950.012
MFVT Sampling
0.975
0.88011.005
0.732
0.867
0. 0.730
MFVI Ours
0.975
0.880
0.73419.000
0.867
2.727
LA Sampling
0.170
0.44410.012
1.238.0 oos
0.46120.008
1.208
LA GLM
0.0
014
1.197
0.
1.18020017
0.086-00.00
0.456+00012
1.068:.055
0.353:0.012
1.284400048
MFVI Sampling
0.133:000m
0.641-0.022
1.09110055
101000.01
1.577
MFVI Ours
0.087
0.468190112
1 1.00
0.619
1. 1.234
LA Sampling
0.006
0.022
0.197
0.129
0.070
LA GLM
0.011
0.024
0.155
0.053
0.057
ECE L
LA Ours
0.009
0.027
0.042
0.017
0,130
MFVI Sampling
0.015
0.070
0.075
0.079
0.118
Ours
Table 5: Performance on language understanding tasks.
Metrics
Methods
BoOLQ
MRPC
LA Sampling
0.383=0000
0,300
0.336#9.011
ACC1
LA GLM
0.698+0.00s
0.611
0. 0.710
LA Ours
0.641 06006
0.500
0.665
LAA. Sampling
0.814_0.006
0.858 021
1.21920.018
NLPD
GLM
0.650
7..000..mm 027
660
LA Ours
0. 684--000
0. 719.0.010
O. 62920.009
LA Sampling
0.259
0.261
0.484
ECE
GLM
0.126
0.1415
LA Ours
0.129
0.118
0.054
the SVHN data (Netzer et al., 2011). Fig. 4 shows the kernel density of the predictive entropy computed on the test sets of CIFAR-10 and SVHN, where the model should have high entropy for data different from the training data. Although our method is slightly underconfident on the in-distribution data, the entropy for in-distribution and OOD data is clearly separated, especially for MFVI.
4.3 CAN OUR METHOD ESTIMATE INPUT SENSITIVIES?
We demonstrate that our method can estimate sensitivities w.r.t. the imputs to the network. For this, we use 3-class MLP trained on the digits 0/6/8. 6.. Our goal is estimate sensitivity maps by assuming that the input images N (xx Σ) are distributed according to a Gaussian centred at the pixel values with diagonal covariance E. We optimise the input covariance of each image by minimising the loss
Σ> n= cross-entropy (f (#x,), yn) H(N(w,, E..)
(11)
In words, we jointly minimise the cross-entropy loss, after analytically propagating the input distri- bution through the network, while maximising the entropy H (xxx,, Ex)) of the input distribution. The optimisation is stopped once the difference in NLPD between the current iteration and initial condition is more than 0.1. Fig. 5 shows examples of the resulting sensitivity maps for deterministic MLP (MAP) and the same MLP with last-layer LA (Bayes). We observe that the largest sensitivity for the digits 0 and are generally in the middle, while for 6 in the upper right corner. The Bayes model shows less spurious sensitivities across the pixels compared to the MAP model. Thus, indicating that incorporating all sources uncertainties can lead to a more interpretable sensitivity analysis.
5 DISCUSSION & CONCLUSION
In this work, we proposed to streamline prediction in Bayesian deep learning through local lineari- sation and local Gaussian approximations of the network. For this, we discussed the propagation in different neural network architectures and covariance structures. In particular, we discussed how to handle Kromeckkerfaactrrised posterior covariances and transformer architectures. We showed through series of experiments that method obtains high predictive performance, provides useful predictive uncertainties, and can be used for sensitivity analysis. Our method helps to make BDL more useful in practice and expands the cases and sources of uncertainties that can be considered.

---

Published as a conference paper at ICLR 2025
LA MAP
LA GLM
LA Ours
MFVI Sampling
MFVI Ours
Density
Entropy
Entropy
Entropy
Entropy
Entropy
Figure 4: Kernel density plots over the predictive entropy from ViT network finetuned on CIFAR-10 10 (blue, in-distribution) and data from SVHN (red, out-of-distribttion). Our method results in clear separation between the in- and
data.
y
y=
y =
6
y =
y =
y
8
Bayes MAP
6
8
8
6
6
8
8
Figure 5: Pixel sensitivity maps of an MLP trained on a subset of MNIST digits (classes 0/6/8). The rows show sensitivities to pixel perturbations for the MLP (MAP) and MLP with last-layer Laplace approximation (Bayes) respectively. The sensitivities are visualised in the range (0.5 1.0). The Bayes MLP shows less spurious sensitivities across the pixels compared to MAP
In future work, we aim to apply our approach to tasks with a larger number of output classes, explore additional use-case scenarios in which streamlined approach can be beneficial, and scale larger networks. Moreover, we aim to investigate further the computational benefits obtained by exploiting the posterior covariance structure and sparsity in the network.
The local linearisation of activation functions induces an error that depends on both the activation function and the location and scale of the distribution over the input to the activation function. Moreover, we assume independence between the activations and model parameters for the local Gaussian approximation in linear layers and residual connections, which may incur loss of information in the propagation. Especially, the independence assumption in the residual block is potentially harmful, and relaxing would be valuable future direction. Further, it would be interesting to estimate the induced approximation eтror to identify potential failure modes. Finally, we assume access to a validation set for fitting the scaling factor of the predictive posterior distribution, which is currently done using grid search. An interesting future step is the use of the marginal likelihood to optimise the scaling factor.
Limitations
ACKNOWLEDGMENTS
AS and RL acknowledge funding from the Research Council of Finland (grant number 339730, 362408). MT acknowledges funding from the Research Council of Finland (grant number 347279). MK acknowledges funding from the Finnish Center for Artificial Intelligence (FCAI). We acknowi- edge CSC - IT Center for Science, Finland, for awarding this project access to the LUMI supercom- puter, owned by the EuroHPC Joint Undertaking, hosted by CSC (Finland) and the LUMI consortium through CSC We acknowledge the computational resources provided by the Aalto Science-11 IIT project. Lastly, we thank Jonas Vestergaard for finding a bug in our code and his feedback on the manuseript.
REFERENCES
Anton Baumann, Rui Li, Marcus Klasson, Santeri Mentu, Shyamgopal Karthik, Zeynep Akata, Amo Solin, and Martin Trapp. Post-hoc probabilistic vision- lssionlanguage models. arXiv preprint arXiv:2412.0601. 2024." 2
10

---

Published as a conference paper at ICLR 2025
Edmon Begoli, Tanmoy Bhattacharya, and Dimitri Kusnezov. The need for uncertainty quantification in machine-assisted medical decision making. Nature Machine Intelligence, 1(1):20 23, 2019.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the Americun Statistical Association, 112(518):59-877, 2017. 2.7
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In Proceedings of the 32th International Conference on Machine Learning (ICML), volume 37 of Proceedings of Machine Learning Research, pp. 1613-1622. PMLR, 2015. 1,2
Javier Burroni, Justin Domke, and Daniel Sheldon. Sample average approximation for black- box variational inference. In Proceedings of the 40th Conference on Uncertainty Artificial Intelligence (UAI). AUAI Press, 2024. 2
Gong Cheng, Junwei Han, and Xiaoqiang Lu. Remote sensing image scene classification: Benchmark and state of the art Proceedings of the IEEE, (05(10):1865: 1800::8888-8888 1883, 2017. 7
Mircea Cimpoi, Subhransu Maji. lasonas Kokkinos, Sammy Mohamed, and Andrea Vedaldi. De- scribing textures in the wild. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 3606-3613, IEEE Computer Society, 2014.
Beau Coker, Wessel P Bruinsma, David R Burt, Weiwei Pan, and Finale Doshi- Wide mean- field bayesian neural networks ignore the data. In Proceedings of the twenty fifth International Conference on Artificial Intelligence and Statistics (AISTATS), volume 131 of Proceedings of Machine Learning Research, pp. 5276-5333 PMLR, 2022. 1
Erik Daxberger, Agustimus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace redux effortless Bayesian deep learning. In Advances in Neural Information Processing Systems 34 (NeurIPS), volume 34, pp. 20089-20103, Curran Associates, Inc., 2021a. 2, 3, 30
Erik Daxherger, Eric Nalisnick, James U Allingham, Javier Antorán, and José Miguel Hernández- Lobato. Bayesian deep learning via subnetwork inference. In Proceedings of the 38th International Conference on Machine Learning (ICML), volume 119 of Proceedings of Machine Learning Research, pp. 2510 25211 PMLR, 2021b. 2,
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 248-255. IEEE Computer Society, 2009. 7
Nikita Dhawan, Sicong Huang, Juhan Bae, and Roger Baker Grosse Efficient parametric approx- imations of neural network function space distance. In Proceedings of the 40th International Conference on Machine Learning (ICML), Proceedings of Machine Learning Research, pp. 7795- 7812. PMLR, 2023. 3
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition scale. In International Conference on Learning Representations (ICLR), 2021. 7
Runa Eschenhagen, Erik Daxberger, Philipp Hennig, and Agustinus Kristiadi. Mixtures of Lapkace approximations for improved post-hoc uncertainty in deep learning. In NewrIPS workshop on Bayesian deep learning, 2021. 2
Andrew Foong, David Burt, Yingzhen Li, and Richard Turner. On the expressiveness of approximate inference in Bayesian neural networks. In Advances in Neural Information Processing Systems 33 (NeurIPS), pp. 15897-15908. Curran Associates, Inc., 2020.
Vincent Fortuin, Adrià Garriga-Alonso Sebastian W Ober, Florian Wenzel, Gunnar Rätsch, Richard E Turner. Mark van der Wilk, and Laurence Aitchison. Bayesian neural network priors revisited. In International Conference Learning Representations (ICIR), 2021. I
11

---

Published as a conference paper at ICLR 2025
Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In Proceedings of the 33th International Conference on Machine Learning (ICML), volume 48 of Proceedings of Machine Learning Research, pp. 1050-1059. PMLR, 2016. 1,2
Yarin Gal, Riashat Islam, and Zoubin Ghahramani. Deep bayesian active learning with image data. In Proceedings of the 34th International Conference on Machine Learning (ICML), Proceedings of Machine Learning Research, pp. 118--1192. PMLR, 2017. 3
Yoav Gelberg, Tycho F. A. van der Ouderaa, Mark van der Wilk, and Yarin Gal. Variational inference failures under model symmetries: Permutation invariant posteriors for Bayesian neural networks. In ICML 2024 Workshop on Geometry grouudded Representation Learning and Generative Modeling, 2024. 1
Ryan Giordano, Martin Ingram, and Tamara Broderick. Black box variational inference with deterministic objective: Faster, more accurate, and even more black box. Journal of Machine Learning Research, 25(18):1-33 2024. 2
James-A Goulet, Luong Ha Nguyen, and Saeid Amiri. Tractable approximate Gaussian inference for
Bayesian neural networks. Journal of Machine Learning Research, 22(251):1-23, 2021. 2
Marton Havasi, Rodolphe Jenatton, Stanislav Fort, Jeremiah Zhe Liu, Jasper Snoek, Balaji Lakshmi- narayanan, Andrew Mingbo Dai, and Dustin Tran. Training independent subnetworks for robust prediction. In International Conference on Learning Representations (ICLR), 2021. 2
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp- 770-778. IEEE Computer Society, 2016. 2
Dan Hendrycks, Kevin Zhao, Sebastian Basart, Jacob Steinhardt, and Dawn Song. The many faces of robustness: A critical analysis of diito--distrbuution generalization. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 8340 8349. IEEE, 2021. 7
Alexander Immer, Matthias Bauer, Vincent Fortuin, Gunnar Rätsch, and Khan Mohammad Emtiyaz. Scalable marginal likelihood estimation for model selection in deep learning. In Proceedings of the 38th International Conference on Machine Learning (ICML), Proceedings of Machine Learning Research, pp. 4563-4573, PMLR, 2021a. 3
Alexander Immer, Maciej Korzepa, and Matthias Bauer. Improving predictions of Bayesian neural nets via local linearization. In Proceedings of the twenty forth International Conference on Artificial Intelligence and Statistics (AISTATS), volume 130 of Proceedings of Machine Learning Research, pp. 703-711. PMLR, 2021b. 2, 3.7
Peter JT Kampen, Gustav RS Als, and Michael Ris Andersen. Towards scalable Bayesian transform- ers: Investigating stochastic subset selection for nlp. In Proceedings of the 40th Conference on Uncertainty in Artificial Intelligence (UAI). AUAI Press, 2024. 2
Raymond Kan. From moments of sum to moments of product. Journal of Multivariate Analysis, 99
(3):542-554, 2008 18
Markelle Kelly, Rachel Longjohn, and Kolby Nottingham. The UCI machine learning repository, 2023. URL: httpss//arccive.is.u....o 6
Agustinus Kristiadi, Matthias Hein, and Philipp Hennig. Being bayesian, even just bit, fixes overconfidence in relu networks. In Proceedings of the 37th International Conference on Machine Learning (ICML), volume 119 of Proceedings of Machine Learning Research, pp. 5436-5446. PMLR, 2020. 2,5
Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, and Vincent Fortuin. Promises and pitfalls of the linearized Lapkace in Bayesian optimization. In Fifth Symposium on Advances in Approximate Bayesian Inference, 2023.
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical report, Toronto, ON, Canada, 2009.
12

---

Published as a conference paper at ICLR 2025
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in Neural Information Processing Systems 30 (NeurIPS), volume 30, pp. 1402-6413 Curran Associates, Inc., 2017. 2
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient- based learning applied to document recognition. Proceedings of the IEEE, 86( (1:2218:222222 2224,, 1998. 6
Jiachen Lu, Jinghan Yao, Junge Zhang. Xiatian Zhu, Hang Xu, Weiguo Gao, Chunjing Xu, Tao Xiang, and Li Zhang. Soft: Softmax-free free transfomerr with linear complexity. In Advances in Neural Information Processing Systems 34 (NeurIPS), pp. 21297-21300 Curran Associates, Inc., 2021.
David JC MacKay, Information- based objective functions for active data selection. Neural Computa- tion, 4(4):590-604, 1992a. 2,
David JC MacKay. Bayesian interpolation. Neural computation, 4(3):415 1992b. 5
David JC MacKay. Bayesian methods for backpropagation networks. In Models of Neural Networks III: Association, Generalization, and Representution, pp. 211-254. Springer, 1996. 3, 7
Wesley Maddox, Pavel Izmailov, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. A simple baseline for Bayesian uncertainty in deep learning. In Advances in Neural Information Processing Systems 32 (NeurIPS), pp. 13132-13143. Curran Associates, Inc., 2019.
James Martens and Roger Grosse. Optimizing neural networks with fronecker-factored approximate curvature. Proceedings of the 32nd International Conference on Machine Learning (ICML), Proceedings of Machine Learning Research, pp. 2408-2417 PMLR, 2015. 6
Lassi Meronen, Martin Trapp, and Amo Solin. Periodic activation functions induce stationarity. In Advances in Neural Information Processing Systems 34 (NeurIPS), pp. 1673-1685. Curran Associates, Inc., 2021.
Lassi Meronen, Martin Trapp, Andrea Pilzer, Le Yang, and Amo Solin. Fixing overconfidence in dynamic neural networks. In Proceedings of the IEEF/CVF Winter Conference on Applications of Computer Vision (WACV), pp. 2680-2690, 2024. 2
Rhiannon Michelmore, Matthew Wicker, Luca Laurenti, Luca Cardelli, Yarin Gal, and Marta Kwiatkowska. Uncertainty quantification with statistical guarantees in end-to-end autonomous driving control. In IEEE International Conference on Robotics and Automation (ICRA), pp- 7344-7350. IEEE, 2020.
Saralees Nadarajah and Tibor K Pogány. On the distribution of the product of correlated normal random variables. Comptes Rendus. Mathématique, 354(2):201-204, 2016. 18
Eric Thomas Nalisnick, On priors for Bayesian neural networks. University of California, Irvine, 2018. 1
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Baolin Wu, Andrew Y Ng, et al. Reading digits in natural images with unsupervised feature learning. In NIPS workshop on deep learning and unsupervised feature feature learning, 2011.
Theodore Papamarkou, Maria Skoularidou, Konstantina Palla, Laurence Aitchison, Julyan Arbel, David B. Dunson, Maurizio Filippone, Vincent Fortuin, Philipp Hennig, José Miguel Hernández- Lobato, Aliaksandr Hubin, Alexander Immer, Theofanis Karaletsos, Mohammad Emtiyaz Khan, Agustinus Kristiadi, Yingzhen Li, Stephan Mandt, Christopher Nemeth, Michael Osborne, Tim G. Rudner, David Rügamer, Yee Whye Teh, Max Welling, Andrew Gordon Wilson, and Ruqi Zhang. Position: Bayesian deep learning is needed in the age of large-scale ai. In Proceedings of the 41st International Conference on Machine Learning (ICML), volume 235 of Proceedings of Machine Learning Research. PMLR, 2024.
Felix Petersen, Aashwin Ananda Mishra, Hilde Kuehne, Christian Borgelt, Oliver Deussen, and Mikhail Yurochkin. Uncertainty quantification via stable distribution propagation. In International Conference Learning Representations (ICLR), 2024. 4,
13

---

Published as a conference paper at ICLR 2025
Apostolos F Psaros, Xuhui Meng. Zongren Zou, Ling Guo, and George Em Karniadakis. Uncertainty quantification in scientific machine learning: Methods, metrics, and comparisons. Journal of Computational Physics, 477:1 2023.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenAl blog, 2019.
Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable Laplace approximation for neural networks. In International Conference on Learning Representations (ICLR), 2018. 6, 7
Simo Särkkä and Lennart Svensson. Bayesian Filtering and Smoothing. Cambridge University Press, 2023. 4
Aidan Scannell, Riccardo Mereu, Paul Edmund Chang, Ella Tamir, Joni Pajarinen, and Amo Solin.
Function-space parameterization of neural networks for sequential learning. In International Conference on Learning Representations (ICLR), 2024. 2
Yuesong Shen, Nico Daheim, Bai Cong, Peter Nickl, Gian Maria Marconi, Bazan Clement Emile Mar- cel Raoul, Rio Yokota, Iryna Gurevych, Daniel Cremers, Mohammad Emtiyaz Khan, and Thomas Möllenhoff. Variational learning is effective for large deep networks. In Proceedings of the s1 International Conference on Machine Learning (ICML), volume 235 of Proceedings of Machine Learning Research. PMLR, 2024. 2, 3, 7,
Freddie Bickford Smith, Andreas Kirsch, Sebastian Farquhar, Yarin Gal, Adam Foster, and Tom Rainforth. inddiction-orientedd bayesian active learning. In Proceedings of the twenty sixth Inter- national Conference on Artificial Intelligence and Statistics (AISTATS), Proceedings of Machine Learning Research, pp. 7331-7348. PMLR, 2023. 3
Ba-Hien Tran, Simone Rossi, Dimitrios Milios, and Maurizio Filippone. All you need is good functional prior for Bayesian deep learning. Journal of Machine Learning Research, 23(74):
2022. 1
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, t ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems 30 (NeurIPS). Curran Associates, Inc., 2017. 5, 16, 21
Maxime Vono, Nicolas Dobigeon, and Pierre Chainais. High-dimensional Gaussian sampling: a review and a unifying approach based on stochastic proximal point algorithm. SIAM Review, 64 16, 2022.
Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Superglue: A stickier benchmark for general-purpose language understanding systems. In Advances Neural Information Processing Systems 32 (NeurIPS), pp. 3266-3280. Curran Associates, Inc., 2019a. 7, 25
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. In International Conference on Learning Representations (ICLR), 2019b. 7, 25
Florian Wenzel, Kevin Roth, Bastiaan Veeling, Jakub Swiatkowski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. How good is the Bayes posterior in deep neural networks really? In Proceedings of the 37th International Conference on Machine Learning (ICML), volume 119 of Proceedings of Machine Learning Research, pp. PMLR, 2020.
Andrew Gordon Wilson. The case for Bayesian deep learning, arXiv preprint arXiv:2001.10995 2020. 2
Andrew Gordon Wilson and Pavel Izmailov. Bayesian deep learning and probabilistic perspective of generalization. In Advances in Neural Information Processing Systems 33 (NeurIPS), pp. 4697-4708. Curran Associates, Inc., 2020. 2,
14

---

Published as a conference paper at ICLR 2025
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, and Jamie Brew. Huggingface transformers: State-of-the-art natural language processing. arXiv preprint 21t.:11910.0377,200999
7
Pierre Wolinski and Julyan Arbel. Gaussian pre-activations in neural networks: Myth or reality? arXiv preprint arXiv:2205. 12379, 2022. 4
Anqi Wu, Sebastian Nowozin, Edward Meeds, Richard Turner, José Miguel Hernández-Lobasd and Alexander L. Gaunt. Deterministic variational inference for robust Bayesian neural networks. In International Conference on Learning Representations (ICLR), 2019. 2,
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv: 1708.77708.,7777, 2017. 6
Adam X Yang, Maxime Robeyns, Xi Wang, and Laurence Aitchison. Bayesian low-rank adaptation for large language models. In International Conference on Learning Representations (ICLR), 2024.
Jiancheng Yang, Rui Shi, Donglai Wei, Zeju Liu, Lin Zhao, Bilian Ke, Ziyang Shi, Yunzhu Li, Xiaoyang Hu, Yang Gao, Ye Xu, Daniel L. Rubin, and Holger R. Roth. Medmnist v2: A large-scale lightweight benchmark for 2d and 3d biomedical image classification. Scientific Data, 10(1):1-14, 2023. doi: 11103//41199-022-05555555 02-5222-11
Cheng Zhang, Judith Bütepage, Hedvig Kjellström, and Stephan Mandt. Advances in variational inference. IEEE Transactions on Pattern Analysis and Machine Intelligence, 41(8):2008-200 2018. 2
15

---

Published as a conference paper at ICLR 2025
APPENDICES
The appendices are structured as follows: App. A presents the derivations our method in detail. App. describes the experimental setup and additional experimental results.
A DERIVATIONS
In this section, we derive how to propagate the distribution deterministically. See Table for the list of notations that will be used throughout this section.
We first derive the general result in App. A. where the posterior covariance has a full structure the linear layer and evaluate the quality of local Gaussian approximation in App. A.2. Next, in App. A.3 and App. A.4 we give the derivation for diagonal and KFAC covariance, respectively. Then, App. A.5 shows the derivation for activation functions. Finally, App. A.6 describes how we apply our method transformer network (Vaswani et al., 2017).
Table 6: Notation.
lowercase bolder letter, vector
W
uppercase bold letter, matrix set
D
T
1th element of x
WK
h row, 7th column of W 4h row of a matrix
W :
k, I
dimension of the output dimension of the imput data feature dimension number of data points total number of classes layer index
i
d
.N
C
A.1 DERIVATION FOR FULL COVARIANCE STRUCTURE
Denote the weight and bias of the mh linear layer as W(m) E D and (m) e Ran respectively, and its input as a (m-1) Dr, The pre-activation then given as (m) (mm) (m-1) + (m) with its thth element being hmm = E2n W (m)m) (m-1) + (m)
We make the following assumptions to obtain tractable distribution on the pre-activation:
Assumption 1: We assume each a; (m-11 is Gaussian distribution.
Assumption 2: We assume that the activations of the previous layer a." (m-1) and parameters the th layer are independent.
From assumption because a; ppn-))((m)) and (((²)) are all Gaussian distributions, A((")) will follow Gaussian distribution as well. We call this local Gaussian approximation as we approximate (m) is Gaussian, (m) will be each local component a (m- with Gaussian. As each h(())
w-fmm
jointly Gaussian. We derive its mean and covariance and drop the layer index if is clear from the context.
16

---

Published as a conference paper at ICLR 2025
Derivation of mean
As as is assumed to be uncorrected with Wkis we have
Da
B [14]] = E
Wa:a; + bn
(12)
EE [Wksa; bk]
(13)
Dn
E [Wssass] + 18 [bx]]
(14)
i=1
Du
E E E [a.] + [bb]
(Assumption 2)
Derivation of covariance
The covariance between the 1tth and 1th pre-activation can be written as
Dn
Dr
Cov [hy, hu] = Cov
04 Wki + bk.
Σ
a; Wi: + bj
(15)
Di
Cov
0g W
a Wi
+ Cov
Σ Whisby bi + Cov Σ ag Wis bb
+ Cov [bx;
(16)
Σ
Cov [a a, Wiyl +
Σ
(Cov [a by] - Cov (i,
DS
1SiSD.
+ Cov [bx, bd]
(17)
We first derive the form of Covla,Wis, a.]]:: Cov (a, Wks a, Wiy]
= [(a B (a,
E Wi,ll
(18)
= E (a,Waaa Wij a. W a,,,,,,,WBBBlpppp...
la Bl [0, Wi,ll
(19)
= E [ara, WR Wiy1 E [a, E [ayWW,
a, W E [ay + E [a, E [ay Wisl
(20)
E. [ E W ke Wiy1
E[a W
[ay]
[Wi
(Assumption 2)
E [ay] Cov (E [a
(as, (3))))))))
W
E [iy]] Cov [Wae, Wig))
E |a
E W
E [ay] E (Wi))
(21)
= E [a.] E [ay] Cov [Wki, Wip + E [a.] E [ig] Cov (ai, ag} + Cov [as, az] Cov [Wki, Wiyl
(22)
Then derive the form of Covla,Wain bij:
Cov (a, ke, b.] = E [[a, Wt - Wa,l)) - [b.))
(23)
IE [[a B [a.] (b: E [6:]]
(Assumption 2)
E (a,b
Wt.
[bi]
[a.] [ E [br]]
(24)
= EE
[a, Wa
E
[a.]
E
[
[by]
(25)
z E
[a.]
E W
by
[
E
W
[b:]
(Assumption 2)
E
[a.]
E W
E
[bx]
Cov Wes, be]) E [as] [Wed] (by)
(26)
E
[a.]
Cov Was by]
(27)
17

---

Published as a conference paper at ICLR 2025
Putting it together, we have Coviha hi]
Din
Cov W,y]
BB [aj] Cov [Wki, bi] + E [a.] Cov [Wi, bx]) + Cov (bk, bi, (28)
15t.SDm
where Cov Novla;WWis a2 Wiy]
E [a [ay] Cov [Wki, Wiy] + [ [Wyy] Cov (a,, az} + Cov (a,, a,} Cov [Wis, Wiyl (29)
Note that EisIJED, Cov(a, Wks aj Wiyl in Eq. (28) could be rewritten into the form matrix multiplication for efficient implementation:
Σ
Cov (a, Wiyl
(30)
Istgs Pa
[[o. Cov (Wa. Wiy + [[ [WWs) Cov [ee: Cov [g
Cov Cov (Wass Wyl
E E|a Cov(War. Whll
[o:]] [an,] Cov[Wa1, Wingl
Σ
[an..
Clv Wol
E [a: [aa]E[[a,,) Cov Weng
Cov[Wai, Wil
Cov[Wei. Win
(33)
Cov(WeD W Wil
Cov(WED Wio
B [[
Waal E [Wina]
Covlas.ai
Covlai, and
+
o
(34)
3 [ B[Wn]
E [WoD Covlas, a
B [Winal
Covlan 0il
Cov Covlo DDal
Cov|Wen, Wiil
Cov Wen, Wind
Σ
(35)
Covlap, onl
CovleD ap..l
CovlWED, Will
CovlWED, WiDa1
A.2 ERROR INDUCED THROUGH LOCAL GAUSSIAN APPROXIMATION
In this section, analyse the error induced by the local Gaussian approximation. Recall that made these two assumptions for the derivation:
Assumption We assume a: (m- 1) 1mm-1) w-(m) is a Gaussian distribution.
Assumption 2: We assume that the activations of the previous layer a; (m-11 and parameters of the mh layer are independent.
We first examine the error induced by A2 on the moments for aj Wks. Given two correlated univariate Gaussian D1 and F2, with the joint being
Cov[r1
1):
N
(36)
E
Cov [2i, m2)
from Nadarajah & Pogány (2016); Kan (2008), although the distribution form of 2y is no longer Gaussian and intractable, its mean and variance can be computed analytically as
[11æg] = E [21] E [x2] + Cov [2),
(37)
Var [_] oi + of [x2] 2 + [21 (070 +22 [x1 [x2]) Cov [x1, x2]
(38)
Applying the above result in our case, we have
E (a Waa] E. [a,] [Wa) + Cov (a:, Wkx],
(39)
Var [a, Wail = W.. + + E. + oi E [a,]
(2E L [ + oā, a...) Cov (ai, Wol.
(40)
18

---

Published as a conference paper at ICLR 2025
As Covlo, Wal is intractable, in A2 ignore the correlation between a, and Was which results in E &a x [a,] W C.r Wal.
(41)
Var la, Wel oa, wwn o Wail ww. Ba
(2E [aj] Wad a. Cov (as, Wad]-
(42)
Note that in the case of diagonal posterior covariance, as each parameter is independent of each other, A2 holds automatically. In this case, we recover the correct mean and variance for a, ks-
Now, we examine the error induced by Al and A2 through Monte Carlo estimation. Fig. 6 provides simulation result illustrating the error induced by the local Gaussian approximation plot the results for weights with the largest absolute magnitude of an MLP trained on MNIST. This approximation works well in practice but fails capture the potential skewness of the distributions.
Wik We
pre aetivativaioon
pre activation
pre pre-activation value
pre pre-activation
pre pre-activation
Figure 6: Comparison between Monte-Cate estimates of the distribution over a, Was and our analytic Gaussian approximation
A.3 DERIVATION FOR DIAGONAL COVARIANCE STRUCTURE
When the posterior has diagonal covariance, the mean E [hx] will still be the same. For covariance, when l: * we have Cov(ha, lu]
Do
ECor Covlo, Wiyl +
Σ
(E [aj] Cov [WAs, bi] + E [a,] Cov [Wi, by]) + Cov [bg,
(43)
Σ Cov (a, W Wiz]
(44)
= Σ
E [a [a Cov [We,, Wi,] + E [Wa] [Wiy) Cov (ai, az] + Cov (as, az) Cov (Wks, Wyj
15t.5DD.
(45)
E [ E Wiy] Cov (as. 2z]-
(46)
D
For h: = we have Var[hy]
Du
Cor a, +
Cov Σ (E [a.] Cov [Wii, bu] + [a.] Cov [War, + Var [bx]
(47)
IS.JED,
(a, Cov a, We.] Var [bx]
Σ Cov Cov
(48)
Dh
Σ [a,] War [Wad + [Was] Var [a.] + Var [a.] Var [Wad + Var [bx]
(49)
Dm
19

---

Published as a conference paper at ICLR 2025
Note that
Var
(m)
I2
4
(m
1)12
Var
W (m)
+
W
[m][m]]
Var (a."
[amm-1))
(50)
15.5
Var (mm-1))
(4 aam Var [
+ Var (b(mm))
(51)
DS
the variance of lig will only rely on the variance of the activations of previous layers, i.e, Varla,m-1) In the case of element-wisg activation functions, Var(a," will only rely Var h (m as now the Jacobian of the activation is diagonal. As a result, in the case where we only need the variance of the input, could drop the computation of Cov [hg, hu] and only compute the variance for each layer, which will largely reduce the computation cost.
A.4 DERIVATION FOR KRONECKER COVARIANCE STRUCTURE
In KFAC, the Hessian is represented in Kronecker product form Hess = A@ B. Denote the prior precision X², then the posterior covariance
Σ (Hess 1T (A & B + X-1
(52)
As there is no closed form for the inverse, to express the covariance in the form of the Kronecker product as well, we approximate the covariance as
Σ (A B A²T)-
(53)
(UAA (U X
(Eigen Decomposition)
[(UAAA + AAUU + ABB)
(54)
(UA(AA+ & (UB(AA AB))
1 A-1 B-1)
Recall for an efficient implementation for computing EistseD, Covla, Wi, a Wiyl (Eq. (35)), we need to retrieve the covariance between the 4thh row of weight and 1th row of weight, which is Dm x Dm matrix:
CovlWai, Wa]
Cov(Wen WIDal
Cor [ W[t.:]]
(55)
CovlWara: Wil
CovlWana WIn
As Σ C D where C E Dim Dm and D E Dom Douts the posterior covariance is represented by a total number of Din x Din matrix with size Doat x Dout Retrieving a Din x Din matrix from it is trivial. In the toy example
shown in Fig. 7, for Dm, 3 and D its covariance is represented by a total number of (Dm x Dm,) matrix I, II, IX with shape 2 x 2 (Dout X Dout). To retrieve Cov [[.,]], :2]
matrix W.
need first decide which Kronecker blocks contain it (in this case block II, III. V and VI) and reconstruct these Kronecker blocks. Then, we retrieve Cov W[1, :). WT [2, :]] from the reconstructed blocks.
In general, the retrieval process consists of two steps: (1) identifying the block indices within the Kronecker product matrix that correspond to the required covariance block and (2) extracting the covariance of interests from the constructed block.
Identifying Block Indices We first identify the Kronecker blocks that contain the covariance of interest. This is achieved by calculating the block indexed for C, which is later used to construct Kronecker blocks. Specifically, the start and end positions of the covariance block corresponding to rows & and / can be computed as:
k- Din
row_start
(56)
Dout
20

---

Published as a conference paper at ICLR 2025
Wi1, WasT WT., W23
WI1, Wsq
Wn1. Wsi
Wn. W W
W
W33
Wiz, Wix Wia. W12
W12, Wis
W1a, Wsi
W12. Ws2
W1a,
Wa3
Cor[W]
Wis, WiygWass W1a
W13, W137
1137 Wa1
W13, Way Wa,, Waa
W21, Wi1 W21, W12
W21, Wi3 W21, Ws1
W21, Wa2 W21, W33
W22, Way 122, W12
W22, Wey 22, Waa
W22, W Wagx W2a, Waa
W2s, Wi1 W2s, W12
W2s, W13 Was, Wsi
W2s, Ws2 Was, Was
Figure 7: retrieve the highlighted submatrix [ov[W[,, [2;]]]
:]] of covariance for W E R2x3, we identify the Kronecker blocks that contain the covariance of interest (II, III, V, and VI), explicate those blocks in memory, and then retrieve the relevant submatrix.
(k: + 11 Din
row_end
(57)
Dout
1. Din Dout
col_start
(58)
(2 [1+11] 1) Dm
col_end
(59)
Dout
Then, can construct the Kronecker blocks that contain the covariance of interest by C' frow_ssie row_end, col_start col_end] D.
Once we have C'row_staae starsttrttt row_end, col_start col_end), and as know the covariance we need to retrieve has shape Dm x Dmn we only need to compute the start row
Extract the Covariance
and column index, which can be computed as
select_row retct_roo_startttt__ (k-D)n mod Douts
(60)
select_coll cotet__co_sartttttttteeeeee = (I- Din) mod Dout-
(61)
A.5 DERIVATION FOR ACTIVATION LAYERS
For a = g(h) where h ~ (h; E [h] Σx) and g(.) is the activation function, we use local linearisation approximate the distribution of Specifically, we do a first-order Taylor expansion on g(-) at E [h]:
a = ggh))) 9 [b]) J Js|n=E[n]( E [h]).
(62)
(63)
Given that Gaussian distribution is closed under linear transformation, we have
h N(E [h], En)
(64)
h E [h] ~V(0, E,,
(65)
Jyln=Eth) (h- - E [h]) ~N Jyln= |[= Ei Jnn-
(66)
(E [h]) + Jy Jg|n=E|N)(1
E [h]) ~N(g(E [h]) JJ |n
En, l==
(67)
Naa: (E [h]) Jo n B(p]
Σn, (()))))))))
(68)
approx
A.6 TRANSFORMER BLOCK
There are four components in each transformer block (Vaswani et al., 2017): (2) MLP; (3) layer normalisation; and (4) residual connection. For MLP blocks, the propagation is the same as described above. For layer normalisation and residual connection, as Gaussian distributions closed under linear transformations, 'pushing" distributions through them is straightforward.
(1) multi-heaa attention;
21

---

Published as a conference paper at ICLR 2025
We describe how to push distributions through attention layers below. Note that for computational reasons, we always assume the input has diagonal covariance.
Given input H RTXD where T' is the number of tokens in the imput sequence and D is the dimension of cach token, denote the query, key and value matrices as Wq E DxD, WK E DDD Wy E DDD respectively, the key, query and value in an attention blocks are
Q HWQ: K HWW, HWW,
(69)
and the output of attention block is
QK T VD
Attention H) Softmax
V
(70)
When the input is a distribution, Q. K and V will all be distributions as well. As pushing a distribution over softmax activation requires further approximation, we ignore the distribution over Q and K for computational reasons and compute their value by using the mean of the input:
Q E [Wo]=
H [Wk]
(71)
For V, for simplicity we describe our approximation for a single token h whose value is v Wyh with hhhh element being e2 Σ Wv,, h4. Assuming h is Gaussian, the covariance between the h and the 1 value is
Cov [Uk: vi] = Cov
Wv his Σ Wrsh
(72)
Cov [Wreltes Wrighj]-
(73)
When treating Wy deterministically, we have
D Cov hw, hy]
Cov [ike, m7]
(definition)
D Wv Wry Cov [bs, hy]
(Wy deterministic)
D
Wv, Wv, War [h:]. (ignore correlation between h for computational reason)
When Wv is an isotropic Gaussian, we have
Cov [uk, uz]
Σ
Cov [Wv. he Wv,hi hgg
(74)
ISi.jDD
Σ 15tJDD
E [hy] + Cov (he, hT]) Cov [Wkr, Wiz] + E [Wa] E [Wy) Cov [h., hy]
(assumption A2)
Σ
W Wiyl Cov [he, hy]
E
(Wy is isotropic Gaussian)
ISijED
Σ ISiSD
E
E Wid Var [4x].
(ignore correlation between h for computational reason)
22

---

Published as a conference paper at ICLR 2025
Var [ux]
Σ
Cov [Wvz, h, hz]
(definition)
isijsD
Σ ISUJSD
[h [hy] + Cov (he, hyl) Cov (Waa, Wayl + E [Waa] E [Wiy] Cov (h4, hy]
(assumption A2)
Σ
EE [h+] Var [b.]) Var [Wid] + [W Var [hi].
ISVED
(Wy is isotropic Gaussian)
(75)
Once we have the distribution V. the distribution over Attention H) becomes distribution of linear combination of Gaussian, which is tractable.
Then for multi-head attention, we assume each attention head's output is independent, which allows us to compute the distribution over the final output in tractable form. As we assume the input is isotropie, we only need compute the variance for each dimension.
A.7 CONVOLUTIONAL NEURAL NETWORK
The derivation for convolutional layers is similar to fully connected layers as convolution layers can be considered as a shared weight fully connected layer. We first give the derivation for convolutional layers, then discuss pooling layers in convolutional neural networks.
Denote the pixel value (i,j)) of Cin channel dcey [i, j]. the Cin channel of convolutional kernel corresponding to Cout" output channel as W co. [] and the pixel value at (ks, I) of the Cout" th output We,..c,
channel as he... [k. I). Then, suppose there are Cnn channels in total and the kernel size is K₁, K we can write the convolutional layer as
hew [Ri, 1)
Σ>Σ
Oen [k - 1,1+ 1JWe..: (i,
(76)
Following our assumption that ae, [k + 1,1+3 j 1] is uncorrelated with W Wen.ra [., have
Derivation of mean
Kn K Σ
C₁
[k
(ae, + 1 1,7 +1
[W ao [1)
(77)
j=1
Derivation of covariance
The covariance between pixels of the Caut output channel are given as:
Cov [heam (2), Nea (Ki, 12]]
(78)
Cov
Σ
Gen [R1 + 11,4 + + 1jW eun en (11.11)
(79)
Ch.1
Kw
aen Ik + 22 - 1,12 + 52 1JW cia [2
(80)
Cn
Kn
ΣΣ
Cov
[ac.. [k1 + 1,44 + j1 1]W can [1.
(81)
Oen [k2 + ig 1,12 + j2 1]W conn [22.22]] (82)
23

---

Published as a conference paper at ICLR 2025
Using earlier results from Eq. (22) and the shorthand W Woon,
Cour, Cin. have:
Cov (dm, [k]
11ri1-1,,0 - 1,44 - W (1, (1
1,, 1,12+ 11gp,-/06_1-__1//.//./0//1a-1111 [ (1 J/|WW//,WW2 jall
(83)
= E [0mm,
-
(84)
E 1 11,44-1,,1+
E 10 12 E [ [1. [ J2]]
(85)
= E [4in,1 [k] - - 1] (22
Cov [1 (22
(86)
E [W | [22) Cov (ain, [hr
(87)
Cov (0mm,
[b + - 1,4 -1
Cov [W W(2.. (88)
B ADDITIONAL EXPERIMENTS
In the Appendix, we provide additional details on (i) the regression experiments App. B.1, (ii) the classification experiments App. B.2, and (iii) the image sensitivity experiment App. B.3. In addition, we also present additional experiments on (i) measuring the performance by varying the number of MC samples for the sampling baseline App. B.4, (ii) estimating the degree of local linearity in our method App. B.5, and (iii) comparing the runtime of our method against the baselines App. B.7 to further demonstrate the benefits of our streamlined prediction with local linearisation and local Gaussian approximation.
B. REGRESSION
Table 7 gives the UCI regression data set information and the neural network structure we used. For all neural networks, Error (RMSE). Our method results in matching or better performance compared with sampling and GLM, indicating the effectiveness of our method. Note that as the mean of the posterior prediction of our method is the same as the prediction made by setting the weights of the neural network to be the mean of the posterior, we result in the same prediction as GLM of LA, and hence the same performance.
the ReLU activation function. In Table 8, report the Root Mean Square
Table 7: UCI regression experiment setup.
Shonhand
(n, d)
Network
SERVO
SERVO
(167, 4)
d.50.1
LIVER DISORDERS
LDD
(345, 5)
d-50-1
AUTO MPG
AM
(398, 7)
d-50-1
ESTATE
(414.6)
d.50.1
FOREST FIRES
FF
(517, 12)
d.50-1
INFRARED THERMOGRAPHY TEMPERATURE
ITT
(1020, 33)
d-100-1
CONCRETE COMPRESSIVE STRENGTH
CCS
(1030, 8)
d-100-1
AIRFOIL SELF-NOISE
ASN
(1503, 5)
4-100-1
COMMUNITIES AND CRIME
CAC
(1994, 127)
d-100-1
PARKINSONS TELEMONITORING
19)
COMBINED CYCLE POWER PLANT
CCPP
(9568,
d-50-50-1
B.2 CLASSIFICATION
Table 9 gives the classification data sets information and the neural network structure we used for the MLP experiment. We use ReLU activation for MLP.
OOD Experiments with MIP To test our method on evaluate the MNIST-trained MLP on rotated versions of the test set as shown in Fig. 8. The rotation degree interval 10" from 0 180" We observe that with increasing rotation degree, our method achieves lower NLPD compared to LA MAP and MFVI Sampling while being close compared with LA Sampling and GLM. Also, our method achieves similar NLPD for both LA and MFVI posterior approximations across the rotation degrees. All methods perform par regarding their ACC. In Fig. 9, we show kernel density plots over the predictive entropy of an FMNIST- MLP evaluated on MNIST. Our method can distinguish between in-distribution and OOD data better than the LA MAP and MFVI Sampling. Although our method under fits the in-distribution data, the separation between them is clear for the OOD data.
(OOD) data, we first
24

---

Published as a conference paper at ICLR 2025
Table 8: Root Mean Square Error I on UCI regression data sets. Our method results in better or matching performance compared with sampling and GLM, indicating its effectiveness.
MFVI Con)
Laplace Approximation (Fall Cov.
(n, d)
Sampling
Ours
Sampling
GLM
Ours
SERVO
(167, 4)
0.749
0. 4..40.
1.632s0
0.658ga:141
0.658_0.10
LD
(345, 5)
0.884
0.88 10.272
0.989
0.9
0.415
AM
(398, 7)
0.415::0.15
0.41770.115
0.505 105
0.371 10010
0.377
1:371=0.10 20.103
REV
(414, 6)
0.563 100006
0.56220.005
0.789 130
0.532
0.532
104
FF
(517, 12)
0.874#1.125
0.874#1.12
0.910
0.852=0.792
0.852 70772
(1020, 33)
0.481
0.497:10.008
0.560:10.075
0.507:00.77
0.507=0.072
CCS
(1030,
0.472
0.476.0.100
0.494
0.301
0.301
ASN
(1503. 5)
0.568=0.062
0.560=0.002
0.550=0.060
0.352
0.352=0.065
CAC
(1994, 127)
0.571 1:0.10
0.585:20.022
1.481 =0 167
0.703
0.703=0.0m
PT
(5875, 19)
60.0000007
0.590=n.008
0.179:0.051
0.110=0.076
0.076
CCPP
(9568, 4)
0.241 +0.008
0.241+0.038
0.338=0.041
0.224
0. 22222200000 00.087
Bold Count
8/11
10/11
2/11
11/11
11/11
Table 9: Classification experiment setup.
Data Set Name
(n, d)
Network Strueture
MNIST
(50000, 784)
d-128-64-10
FMNIST
(50000, 784)
d-128-64-10
ORGANCMNIST
(12975, 784)
d-128-64-11
ORGANSMNIST
(13932, 784)
d-128-64-11
LA
LA
LA Sampling
LA Sampling
LA Ours
LA Ours
NLPD
MFVI Suppling
MFVI Sampling
MFVVOOIS
ACC
MFVI Ours
MAP
MAP
D.4
Rotation Degres
Rotation Degree
Figure 8: NLPD and ACC for MNIST-trained MLP rotated versions of the MNIST test set. The rotation degree interval is 10 from 180% Our method achieves similar NLPD for both LA and MFVI posterior approximations.
LA
Ours
MFVT Sampling
IVON Ours
Density
Entropy
Entropy
Entropy
Entropy
Entropy
Figure 9: Kernel density plots over the predictive entropy from an MLP trained on FMNIST (blue, in- distribution) and data from MNIST (red, out-of-distribution)). Our method results a clear separation between the in- and out-of-distribution data.
Our method applied to MI.P ViT In Table 10 we report the results for fine-tuning the MI.ss after the attention layers in the last two transformer blocks in ViT and later treating them Bayesian. We observe that our method achieves better or on par NLPD and ECE compared to the baselines for both LA and MFVI across all data sets while maintaining similar ACC the baselines.
We present results for GPT-2 on tasks from GLUE (Wang et al., 2019b) and SuperGLUE (Wang et al., 2019a) benchmarks. These natural language understanding tasks could be turned into classification tasks with the prompt shown in Table 11. We add a classification layer on top of the encoder and use the embedding of the last token in each input do classification.
25

---

Published as a conference paper at ICLR 2025
Table 10: Performance metrics using ViT with posterior approximation MLPs after the attention layers with the standard error for ACC and NLPD. Our method achieves better NLPD and ECE in general and achieves similar ACC compared the baselines.
Metrics
Methods
CIFAR-10
CIFAR-100
DTD
RESISC
IMAGENET-R
LA Sampling
2.902
0.855=0.001
0.63620.011
0.81220.005
0.589
LA GLM
0.974
0.873=000
0.71410.010
0.88610.04
0.687
ACC1 T
LA Ours
0.976
0.884+0.005
0.716+0.000
0.90880.001
0.713
MFVT Sampling
0.978
0.896:10.00
0.727
0.87020.001
0. 0..031
MFVI Ours
0.978
0.895
0.720
011220012
LA Sampling
0.169-0.001
1.043=0.010
2.035 022
1.30420.011
2.330+0.041
LA GLM
0.089.0.005
01
1.260
1.584
0.088:0.008
0.457+0001
1.078+0.09
0.318:00012
1.33940047
MFVI Sampling
0.124.0.01
0.480-0.018
1.27710.000
1.480
MFVI Ours
0.081 40.006
0.43611001
1.14619.80
0.650
20620.063
LA Sampling
0.078
0.349
0.442
0.431
0.331
LA GLM
0.005
0.097
0.174
0.157
0.155
ECE L
LA Ours
0.007
0.032
0.055
0.017
0.087
MFVI Sampling
0.014
0.040
0.083
0.075
0.115
Ours
Table 11: Prompt templates for fine-tuning GPT-2 on natural language understanding tasks.
Task
Prompt
MRPC
Answer whether sentence 2 is equivalent sentence 1. Sentence [sentencel] Sentence 2: (sentence2). (senteeeeessseeeeeeee
Answer:
WiC
Select whether word [word] has the same meaning in these two sentences. Sentence I: (sentencel Sentence 2: [sentence2). Answer:
BoolQ
Answer the question with only True or False.
Passage: {passage) ). Question: fquestion) Answer:
Lasy Layer Laplace Approximation VIT In Table 12 report the results for fine-tuning only the last classification layer ViT base and later treating it Bayesian. We observe that our method (LL- LA Ours) achieves better par NLPD and ECE compared to last layer Laplace approximation (LL-LA GLM/Sampling) across all data sets while maintaining similar ACC. Compared to the where more layers are treated Bayesian (L.A Ours) (results taken from Table 4), last layer approximations in general have lower accuracies and higher NL.D and ECE, indicating the henefits gained by treating more layers Bayesian. Table 13 we report the wall-clock run times for last layer Laplace approximation on CIFAR-10 in milliseconds (see App. B.7 for the run time setting) Our method has matching speed with MAP and slight speed improvements over GLM.
Table 12: Performance metrics using ViT with posterior approximation on last layer with the standard eтrо for ACC and NLPD. In the last layer Laplace approximation (LL-LA), our method achieves better NL.D and ECE in general and achieves similar ACC compared to the baselines. Compared with the case where more intermediate layers are treated Bayesian (I.A Ours), last layer Laplace approximation in general has lower accuracies and higher NLPD and ECE.
Meties
Methods
CIFAR-10
CIFAR 100
DTD
RESISC
IMAGENET R
LL-LA GLM
0. 26.000
0.825=0004
0.681=0.00
0.50620011
0.592
1
LL-LA Sampling
0.968=0.002
0.827:0.00
0.604
LI-LA Qurs
0 0.965200m
0.826_0001
0.693 one
0.508
0.592
T.A. Ours
0.97649.0uz
0.88019.005
0.719g0.010
0.802g0.001
0.739200012
GLM
0. 200
O. 0.88920.0
1.500a0.01
2.574gu.00
2.034gm.051
NLPD
LL-LA Sampling
0.18==0.00
2.456=0.0a0
LL-LA Ours
0. 1.00000000
0.874=0001
1.411 100000o00 coe
2.319:00002
2.020 a0.a52
LA Ours
0.08644.006
0.456+0.012
1.068
0.353
1.264:0.06
L.I.AA GLM
0.011
0.062
0.137
0.333
0.145
ECE L
LL-LA Sampling
0.015
0.045
0.215
0.304
0.127
LL-LA Ours
0.007
0.029
0.026
0.233
0.135
LA Ours
0.009
0.027
0.042
0.017
0.130
26

---

Published as a conference paper at ICLR 2025
Table 13: Wallclock times for Last layer ViT base on CIFAR-10 in milliscconds.
Model
Methods
AvG. RUNTIME (#: STD)
MAP
3.732
LA Sampling
088.7324000
Layer
LA GLM
5.517+0.022
Ours
3.782+0.068
y ==0
y 0
y
0
y
y
y
y
== 8
y 8
==8
MAP
6
8
8
8
Bayes
6
8
8
8
Figure 10: Pixel sensitivity of MLP classifiers trained on binary classification tasks (0/6/8) for MNIST digits. The rows show the sensitivity of the MAP predictor pixel perturbations, and the pixel sensitivity for a last-layer Laplace approximation. The predictive distribution is approximated analytically in both cases. We observe that the Bayesian model using Laplace approximation has less spurious sensitivities to pixel perturbations indicating that it is more robust to input perturbations. 1.0)
The sensitivities are visualised in the range (0.5
B.3 IMAGE PIXEL SENSITIVITY
We trained 4 layer MLP classifier on MNIST digits zero and eight using a batch size of 64, learning 5, and for 50 epochs. We used a subset 0.1% of the
rate of le 3, weight decay set training data as held-out validation set and assumed full covariance Gaussian distribution for each input centred the pixel values of the datum and with fixed covariance of 1e 5. Furthermore. we then computed the pixel sensitivities for the trained model by learning the pixel-wise input covariance matrices by minimizing the negative ligelikel of the held-out validation set and jointly maximizing the entropy of the input distributions. The optimisation image independently and using Adam with learning rate of 5e 3 until the validation loss dropped below divergence the initial loss of le 2. Doing typically took around 700 iterations. Fig. 10 shows some additional examples with the input-dependent sensitivities.
performed for each
B.4 EFFECT OF THE NUMBER OF MC SAMPLES ON PERFORMANCE
We investigate the influence of number of samples on performance.
On regression tasks with small scale neural network (two layer MLP), run experiments with the range of [100, 500. 1000, 5000, 10000, 50000]. On classification tasks with medium scale neural net- work (four layer MLP), we run experiments with the range of [100, 500, 1000, 5000, 10000, 25000]. On classification tasks with large scale neural network (ViT-Base), we run experiments with the range of [10, 20, 100). The results are reported Figs. to 13. The number of samples we used to report results in the main paper is shown in the dashed line.
In Fig. 11 and Fig. 13, we observe that for regression on small scale networks and classification on large scale networks, the performance saturates when set the number of samples to 50 samples (classification with ViT-Base) or 1000 samples (regression with two layer MLP). In Fig. 12, the performance saturates with 1000 samples with LA. For MFVI the performance saturates with 5000 samples, as the improvement gained on NLPD is marginal (from 2.13 to 2.11) from 1000 samples to 5000 samples, in experiment set the number of samples 1000 for MFVI well.
27

---

Published as a conference paper at ICLR 2025
MFVI on CIFAR-10 dataset
LA on CIFAR-10 dataset
0.150
0.180
NLPD
0-140
0.175
0.130
0.170
64
100
0.972
0.075
0,970
ACC
0.948
0.974
0.966
1
20
40
60
100
20
40
60
80
100
60
30
Time (min)
15
19
20
20
60
80
100
20
60
80
100
Number of MC samples
Number of MC samples
Figure 11: Effects of the number of MC samples on performance for LA and MFVI on CIFAR-10 with ViT-Base model. In the results reported in main paper, the number of MC samples to 50 (dashed line).
B.5 ESTIMATING DEGREE OF LOCAL LINEARITY
We performed an additional experiment to assess the degree local linearity of trained MLPP with ReI.U activation functions. In particular, for trained MLP f()), we are estimating the expected absolute error
ii = Es-p(x) [=] - (1 + <)]],
(89)
where and òLnn is zero for any if f()) is linear around each
In our experiments we vary e in the range MLP with layers with sizes [784, 128, 64, 10] trained on MNIST digits. After training, removed the softmax operation on the last layer and measured the local linearisation error on the logits. We estimated the inputs and the function outputs on the same subset. The range of input values is 3.246 and the range of the function outputs varies between 153.072 and 291.168. Fig. 14 shows the results for each of the ten output dimensions scaled relative to their respective range. We observe that the trained ReLU MLP obtains low expected absolute error and behaves locally linear certain degree.
E [le-6, le-5,..., 1] for fully connected ReLU
random subset of 124 validation data points and estimated the range of the
B.6 COMPARISON OF VALUE COVARIANCE IN TRANSFORMERS
One way to improve efficiency in transformer is dropping the correlation between values, i.e., drop the correlation Cov tyj given in Eq. (74). We compare the performance of both approximations and the results are given in Tables 14 to 16. Both approximation results almost the same results for NLPD. ACC and ECE.
28

---

Published as a conference paper at ICLR 2025
MFVI on FMNIST dataset
LA on FMNIST dataset
2.200
0.620
NLPD
2.150
0.000
2.100
0.580
10a
10
2.5
102
10a
10 2.5. 10
0.864
0.862
0.833
860
ACC
o
0.858
0.831
0.850
10
10
104
102
10a
10
30 .00000
0100
Time (min)
20 000
4.000
10,000
10000
0,000
0.080
102
103 Number of MC samples
10
2.5 104
10a
10at Number of MC samples
10 2.5 10*
Figure 12: Effects of the number of MC samples performance for LA and MFVI on FMNIST. In the results reported in main paper, we set the number of MC samples to 1000 (dashed line).
Table 14:
Negative Log Predictive Density (NLPD) for ViT with posterior approximation in the attention layers. We compare only considering variance for value V and considering full covariance. For MFVI and LA, both approximations results in almost the same result.
Mean Field Variational Inference
Laplace Approximation
Datasct
Full Covariance
Only Variance
Foll Covariance
Only Variance
C1FAR-100
0.468 = 0.012
0.467 + 0.012
0.456 + 0.012
0.456 + 0.012
DTD
1.006 1 0.035
1.007 1 0.035
1.068 1 0.035
1.068 0.035
RESISC
0.619 0.019
0.616 0.019
0.353 0.012
0.352 0.012
IMAGENET-R
1.234 = 0.052
1.233 + 0.052
1.261 + 0.013
1.267 + 0.013
Table 15: Accuracy (ACC) for ViT with posterior approximation n the attention layers. We compare only considering variance for value V and considering full covariance. For MFVI and LA, both approximations results almost the same result.
Meon Field Variational Inference
Laplace Approximation
Datasct
Full Covariance
Only Variance
Full Covariance
Only Variance
CIFAR- 10
0.975 +: 0.002
0.975 + 0.004
0.976 + 0.000
0.976 0.004
CIFAR-100
= 0.003
+ 0.009
0.880
0.880 0.009
DTD
0.734 = 0.010
0.734 + 0.012
0.719 + 0.010
0.719 0.012
RESISC
0.867 = 0.004
0.867 + 0.009
0.892 +
0.892 0.008
IMAGENET-R
0.727 0.012
+ 0.012
+ 0.012
0.739 0.012
29

---

Published as a conference paper at ICLR 2025
MFVI on CCPP regression dataset
LA on CCPP regression dataset
0.304
1,850
NLPD
0.363
1,800
0.363
1.
0.363
10a
10a
10
10
10a
10a
10a
-00
0.6090
012259
2315
RMSS
0.400
0.237
0 0.200
102
10a
104
5 104
10a
103
104
5- 104
oo0
15,000
6. 000
Time (min)
00000
1,90
3.000
0600
0.000
0.000
102
10
104
204
10a
10
10
10
Number MC samples
Number samples
Figure 13: Effects of the number of MC samples on performance for LA and MFVI on regression tasks. In the results reported in main paper, we set the number of MC samples to 1000 (dashed line).
0.01
0.008
Eppected Absolue Erorr
0.006
0.004
10
10
10a
Value of
Figure 14: Estimated divergence from locally linear function function of Note that value of zero means that the function behaves locally like a linear function.
B.7 RUNTIME EXPERIMENT
We compared the runtime of our method against sampling (using the 'torch-laplace library Daxberger et al. (2021a) for Laplace and the IVON Shen et al. (2024)) and the GLM implementation of the "torch-laplace" library for diagonal posterior covariances. For our streamlined approach on ViT, we assessed two cases: (i) propagating Cov(ex, vi] covariance terms (cf., Eq. (74)) through the transformer (+Cov), and (ii) ignoring Cov [Uk. my] covariance terms. We used pre-trained ViT base model on CIFAR- 10 and pre-trained MLP MNIST. For comparison
also list the runtime for
a single forward pass. For this, we ran experiments on an NVIDIA H100 80GB GPU for 400 data points, batchsize of one, and for each data point repeated the measurement ten times. To account
30

---

Published as a conference paper at ICLR 2025
Table 16: Expected Calibration Error (ECE) for VIT with posterior approximation the attention layers. We compare only considering variance for value V and considering full covariance. For MFVI and LA, both approximations results in almost the same result.
Meon Field Voriational Inference
Laplace Approximation
Dataset
Full Covariance
Only Variance
Full Covariance
Only Variance
CIFAR-10
0.007
0.008
0.009
0.008
CIFAR-100
0.024
0.026
0.027
0.023
DTD
0.040
0.010
0.042
0.043
IMAGENET-R
0.043
0.132
0.130
0.039
for code compilation overheads, dropped the first run each data point. We report the mean and standard deviation of the runtime (in milliseconds) over the remaining nine runs and all 400 data points. The results are shown in Table 17. For VIT, we can see that our method without Cov [1g,, ey] covariance terms has a comparable runtime to single forward pass in the deterministic model. When additionally accounting for covariance terms, overall comparable performance. Note that our implementation is not optimised for speed and larger speedups may be obtained by optimising the code. For MI.,, we obtain slight speed improvements over LA GLM but overall comparable performance.
obtain slight speed improvements over GLM but
Table 17: Wallclock times for VIT base on CIFAR-10 and MLP on MNIST in milliseconds.
Model
Methods
(+:
MAP
3.737 23.700000
LA Sampling
190. 806 137
LA GLM
17.191
MFVI Sampling
207.85111...
Ours (+ Cov)
14.728 144
4.350
MAP
0.069
LÄ Sumpling
D8 98.58423..3
LA GLM
1.656=0.ozz
MFVI Sampling
Ours
0.542 20.073
31


Reviews:
[{"review_id": "kS0PdQB8KX", "reviewer": "Reviewer_8AJz", "summary": "This paper aims to address the practical challenge of efficient uncertainty estimation in Bayesian deep learning (BDL). While recent work has improved posterior estimation in BDL, computing predictions typically requires expensive Monte Carlo sampling.\nThe authors propose a streamlined approach that enables analytical computation of predictions in a single forward pass. Their method combines Local linearisation of activation functions using first-order Taylor expansions and local Gaussian approximations at linear layers.", "questions": "Please see weaknesses above", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper introduces a novel approach to Bayesian Deep Learning (BDL) by eliminating the need for Monte Carlo sampling through local approximations. It nicely blends local linearisation with Gaussian methods and offers interesting solutions for managing transformer architectures and Kronecker-factored covariance.\n\nQuality-wise, the work is solidly validated with extensive experiments across various tasks and architectures. It thoroughly benchmarks against established baselines, clearly showcasing benefits like improved efficiency and reliable uncertainty estimates. Additionally, the method is versatile, supporting different posterior approximations such as Laplace Approximations and MFVI.\n\nThe presentation is clear and well-organised.", "weaknesses": "The treatment of attention layers (using deterministic queries/keys) needs stronger justification. The scaling factor for predictive variance is tuned on validation data - this seems ad hoc. No discussion of potential failure modes or limitations of the local linearisation assumption", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "huKTuFP7zL", "reviewer": "Reviewer_KRRC", "summary": "The paper proposes an alternative to Monte Carlo sampling to compute the predictive distribution of a Bayesian neural network. This is obtained with a single forward pass instead of the multiple ones required by sampling. Some approximations make this possible, namely: Local linearization of activation functions, Gaussian approximation of linear layers, independence assumptions between activations and following layer parameters, deterministic key and query in attention layers and block diagonal structure of multi-head input.\nOn a variety of task, they empirically show superior or matching performances versus LA sampling, GLM and MFVI sampling.", "questions": "Can the author explain why they never mention the number of samples used with MC sampling baseline?\n\nCan the author provide an ablation study on number of samples used?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "The experiment sections cover a variety of tasks.", "weaknesses": "The main claim of the paper is to provide a scalable alternative to sampling, but they fail in providing sufficient evidence for such claim. Not all the experiments details are reported, first of all, they do NOT specify the number of samples used in the MC sampling baseline.\nFor whatever fixed parameter approximate posterior $q$, the proposed method is an approximation of the predictive distribution. Such predictive distribution is exactly what we get in the limit of infinite number of samples with MC sampling. Consequently, MC sampling has to perform better than the proposed method if enough samples are used.\n\nThis is point is not studied and not even addressed. Clearly indicating a malicious behaviour.\nIf the claim of the authors is that their method beats MC sampling \"when not too many samples are used\" then this would make sense, but this is neither the claim of the authors, neither such \"too many\" value is studied or attempted to show.\n\nAnother minor point is that in Line 230-231 the authors write \"residual connection ... resulting distribution can be obtained analytically\". Residual connection is essentially a sum, and the elements involved here are two Gaussians. While it is true that the sum of two Gaussian can be \"obtained analytically\", this step hide a big approximation: you can't store a mixture of Gaussians here, so the only scalable approach is to once again employ independence and \"sum\" the Gaussians erasing most of the information. While such approximation may be ok if it leads to good empirical performance, the way is presented is not clear and tends to hide such approximation step.\n\nTypo and minor details:\n-Line 80 \"recent work (2024) developed ... and have shown good performance (2016) (2019)\". I get what you mean, but the way the sentence is formulated imply that something developed in 2024 was used in 2016 and 2019...\n-Line 260 \"extends our the discussion\" I guess this is a typo\n-Line 324 \"ours results in\"  should be either \"our results in\" or \"ours result in\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Bp19QXYoKk", "reviewer": "Reviewer_oXsH", "summary": "The paper proposes a method for computing an efficient approximation to the predictive distribution given a Bayesian neural network’s (BNN) posterior over the model’s weights. Their method locally linearises the activation functions, allowing for an analytical computation of the predictive distribution assuming a Gaussian weight space posterior approximation. This only requires a single forward pass, in contrast to MC sampling-based methods, and no expensive Jacobians that are necessary when globally linearising a neural network. The application of the method to different architectures and posterior covariance structures is discussed. To demonstrate the applicability of their method, the authors conduct extensive experiments applying their method to Laplace and mean-field variational inference posterior approximations on a range of architectures and datasets.", "questions": "**Minor comments and questions**\n\n- When mentioning stable distributions, please define them informally in one sentence first.\n- Have you considered applying this method to architectures that use convolutions?\n- I find the word “useful” in Figure 1 confusing. What makes this “useful” outlier detection/sensitivity analysis? What is non-useful outlier detection/sensitivity analysis?\n- In the paragraph about ensemble methods in the related work section you mention that they typically require multiple forward passes. However, there are efficient ensembling approaches that only require a single forward pass and could be mentioned here, e.g. MIMO [1] and the combination of MIMO with last-layer Laplace approximations [2].\n- One potential disadvantage of the method is the implementation overhead, as it requires a custom implementation for each layer-type within an architecture. Have you thought about how to improve the usability of your method?\n\n[1] Havasi et al. (2021). [Training independent subnetworks for robust prediction](https://openreview.net/forum?id=OGg9XnKxFAH).\n\n[2] Eschenhagen et al. (2021). [Mixtures of Laplace Approximations for Improved Post-Hoc Uncertainty in Deep Learning](https://arxiv.org/abs/2111.03577).", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Overall, I recommend to accept the paper since it 1) proposes a simple, yet effective solution to the problem of efficient Bayesian neural network prediction, 2) appears to be technically correct and is well presented, and 3) provides extensive empirical evidence for the claim that their method performs at least comparable to more expensive alternatives.\n\nThe proposed method provides an intuitive solution to the problem of efficient BNN prediction, assuming that we want to consider a posterior distribution over more than just the last-layer weights of the neural network. The writing is clear, the figures and tables are nicely done. I also appreciate the detailed derivations in the appendix. While I have some concerns regarding the baselines, the experiments cover significant breadth. In particular, I like that that the sensitivity analysis was included, as this application highlights the more unique strengths of this method (see weaknesses for more on this).", "weaknesses": "The main weakness of the paper is that it does not consider alternative methods for approximating the predictive distribution of a BNN with a single forward pass. You do not compare the method to last-layer variants of the Laplace approximation and MFVI. This is important to determine if there is any benefit in trying to linearise the neural network to allow for more efficient predictions vs. just treating the network up to the last layer deterministically. For example, when using a last-layer Laplace approximation together with a probit approximation (in the case of classification), making a prediction also only requires a single forward pass and no expensive Jacobian computations. However, it is nice that you include results for which the linearisation w.r.t. the inputs is necessary, i.e. the sensitivity analysis. I think this could be even emphasised a bit more by explicitly saying that this is a downside of using last-layer methods.\nStill, in most cases last-layer methods would be sufficient -- in particular also for the larger scale experiments in section 4.2.\n\nFinally, I think the title is too general. As you describe yourself in the introduction, you can deconstruct Bayesian inference into three steps, 1) the specification of the prior, 2) the estimation of the posterior distribution, and 3) the computation of the predictive distribution. Your paper aims to streamline the third step. Hence, a more suitable title appears be “Streamlining prediction in Bayesian deep learning”.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hz1oC96oVR", "reviewer": "Reviewer_tEnH", "summary": "This paper proposes a simple way to obtain the a posteriori distribution over functions represented by Bayesian neural nets (BNNs). The core idea is to assume the product $W^{l} a^{l-1}$ between the weight matrix of layer $l$ and the activations $a^{l-1}$ from the previous layer to be a Gaussian. It follows that the preactivation $h^{l}$ is also Gaussian. Then, by linearizing the activation function $g$, one can approximate the activation vector $a^{l}$ is also Gaussian.\n\nBy induction, the output $a^{L} = f(x)$ is also Gaussian, which can readily be used to do inference in close form. For example, one can apply the probit approx. in classification or used in standard acquisition functions for sequential decision-making.\n\nThe key difference between this approach and previous approaches is the fact that the linearization is done in terms of the activations (and hence the input), and not over the weights as done in the linearized Laplace approximation. Nevertheless, the authors showed that the performance did not suffer from this linearization.\n\nThe important advantage of this approach is its scalability: It only requires a single forward pass with minimal overhead. This is in contrast to MC integration, which requires multiple forward passes, and the linearized Laplace, which requires a backward pass to compute the Jacobian over the weights. Indeed, the authors showed that their method is scalable to vision transformers and (small) language models.", "questions": "Please see the Weaknesses section above. Additionally:\n\n- How plug-and-play is the proposed method? Suppose the user has their model that has been Laplace-approximated via `laplace-torch`. What does the user need to do to enjoy the proposed method?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The proposed method is conceptually simple\n- The proposed method might be (see weaknesses below) more efficient than prior methods in obtaining the output distribution.\n- While linearization $x \\mapsto f(x)$ is performed, the performance doesn't seem to suffer. In fact, it is very close to the linearized Laplace predictive (GLM in the paper).\n- The method seems to be scalable (see weaknesses below).", "weaknesses": "I always appreciate the push for simple methods to make BNNs scalable without sacrificing performance. However, I noticed some issues in the presentation of the paper:\n\n- The idea of linearizing the activation function has been done before in the context of continual learning by [Dhawan et al., ICML 2023](https://proceedings.mlr.press/v202/dhawan23a.html). The authors should cite this work and compare their method in detail. I acknowledge that the purposes are different. However, one cannot ignore the fact that they are very similar.\n\n- It is unclear to me why one can still maintain good performance when the network is effectively linear in its inputs. Intuitively, this should not be the case since the representation learning of the base network is negated by the linearization. Could the authors please elaborate in detail (additional, empirical investigations like measuring the degree of linearity of the resulting linearized $f(x)$ might be one way to go about this).\n\n- The main selling point of the approach is that to obtain the distribution over outputs, only a single pass is needed. However, it's currently unclear how much speedup this method provides compared to, e.g., the linearized Laplace (GLM).  A plot/table comparing the wall-clock times between the proposed method and the GLM predictive would be great. \n\n- To make your selling point stronger, I suggest additional experiments with larger models such as Llama-3 7B etc. Laplace can still be done efficiently over PEFT (see [this](https://proceedings.mlr.press/v235/kristiadi24a.html) and [this](https://openreview.net/forum?id=FJiUyzOF1m&noteId=Uazck8TQ5A)). I believe this functionality is also available in `laplace-torch`: <https://aleximmer.github.io/Laplace/huggingface_example/>. Then, again, please compare the wall-clock times.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)