Variational Bayesian Last Layers
Authors: James Harrison
John Willes
Jasper Snoek
Venue: ICLR 2024 spotlight (2024)
Keywords: bayesian deep learning
variational methods
bayesian last layers
neural linear models
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We introduce a deterministic variational formulation for training Bayesian last layer neural networks that improves accuracy and calibration for free.

Abstract: We introduce a deterministic variational formulation for training Bayesian last layer neural networks. This yields a sampling-free, single-pass model and loss that effectively improves uncertainty estimation. Our variational Bayesian last layer (VBLL) can be trained and evaluated with only quadratic complexity in last layer width, and is thus (nearly) computationally free to add to standard architectures. We experimentally investigate VBLLs, and show that they improve predictive accuracy, calibration, and out of distribution detection over baselines across both regression and classification. Finally, we investigate combining VBLL layers with variational Bayesian feature learning, yielding a lower variance collapsed variational inference method for Bayesian neural networks.

Published as a conference paper at ICLR 2024
VARIATIONAL BAYESIAN LAST LAYERS
James Harrison's .Jhn Willes?, Jasper Snoek
l Google DeepMind, Vector Institute
samesharrii wonggoogleengnnoeeoooooooo com, jon.w..Leesveettolsstttteee jsnoekegoogle.co com
ABSTRACT
We introduce deterministic variational formulation for training Bayesian last layer neural networks. This yields sampling-free single-pass model and loss that effectively improves uncertainty estimation. Our variational Bayesian last layer (VBLL) can be trained and evaluated with only quadratic complexity in last layer width, and is thus (nearly) computationally free to add to standard architectures. We experimentally investigate VBLLs, and show that they improve predictive accuracy, calibration, and out of distribution detection over baselines across both regression and classification. Finally, we investigate combining VBLL layers with variational Bayesian feature learning, yielding lower variance collapsed variational inference method for Bayesian neural networks.
INTRODUCTION
Well-calibrated uncertainty quantification is essential for reliable decision-making with machine learning systems. However, many methods for improving uncertainty quantification in deep learning (including Bayesian methods) have seen limited application due to their relative complexity over standard deep learning. For example, methods such as sampling-based mean field variational inference (Blundell et al., 2015), Markov chain Monte Carlo (MCMC) methods (Papamarkou et al., 2022; Neal, 1995; Izmailov et al., 2021), 1), and comparatively simple heuristics such as Bayesian dropout (Gal & Ghahramani, 2016) all have substantially higher computational cost than baseline networks. Single- pass methods (where only one network evaluation required) often require substantial modifications to network architectures, regularization, or training and evaluation procedures, even for the simplest such models (Liu et al., 2022; Wilson et al., 2016b; Kristiadi et al., 2021).
In this work, we take a simplicity first approach to Bayesian deep learning, and develop a conceptu- ally simple and computationally inexpensive partially Bayesian neural network. In particular, we investigate variational learning of Bayesian last layer (BLL) neural networks. While BLL models consider only the uncertainty over the output layer of the network, they have been shown to perform comparably to more complex Bayesian models (Watson et al., 2021; Harrison et al., 2018; Fiedler & Lucia, 2023; Kristiadi et al., 2020). Our variational formulation relies on deterministic lower bound on the marginal likelihood, which enables highly-efficiens mini-batch, sampling-free loss computation, and is thus highly scalable.
Contributions. Concretely, the contributions of this work are:
We present variational Bayesian last layers (VBI..s), novel last layer neural network component for uncertainty quantification which can be straightforwardly included in standard architectures and training pipelines (including fine-tuning), for both deterministic and Bayesian neural networks. We derive principled and sampling-fres free Bayesian training objectives for VBLLs, and show that with careful parameterization they can be computed the same cost standard training, and trained with standard mini-batch training.
We show that VBLLs improve predictive accuracy, likelihoods, calibration, and out of distribution detection across wide variety of problem settings. We also show VBLLs strongly outperform baseline models in contextual bandits.
We release casy-to-use package providing efficient VBLL implementations in PyTorch.
2 BAYESIAN LAST LAYER NEURAL NETWORKS
We first review Bayesian last layer models which maintain a posterior distribution only for the last layer in a neural network. These models correspond to Bayesian (linear or logistic) regression or Bayesian Gaussian discriminant analysis (for each of the three models present, respectively) with

---

Published as a conference paper at ICLR 2024
learned features. We assume T total data points, and write inputs as E R R.., . For regression, outputs are E R Nv; ; for classification, outputs are E {l..... Nyy and denotes the N, Ny.ddimensional one-hot representation. For all models discussed in this section, we will use neural network features Φ : R N= x + + N These correspond to all parts of a network architecture but the last layer, where 0 E e denotes the weights of the neural network. We will typically write Φ == (1x, O) for notational convenience and refer to these parameters as features because they define the map from inputs the feature embedding which the BLL operates.
2.1 REGRESSION
The canonical BLL model for the regression case is
y = w T $(2 0) +
(I)
where is assumed to be normally distributed with zero mean and covariance E, and these noise terms are i.i.d. across realizations. We specify Gaussian prior* p(w) = =N assumed independent of the noise e. Posterior inference in the BLL model is analytically tractable for a fixed set of features.
The marginal likelihood may be computed either via direct computation or by iterating over the dataset. Fixing a distribution over w of the form ((w, 5). the predictive distribution is
p(y ær n, 0) = N(w T Φ Φ SΦb + E) where 17 denotes the parameters of the distribution, here 17 = (w.
(2)
2.2 DISCRIMINATIVE CLASSIFICATION
In this subsection we introduce BLL model that corresponds to standard classification neural networks, where
p(y æx W, 0) = softmax(z),
x W WΦ(x, 0) +
(3)
where E Nv are the logits. These are also interpreted as unnormalized joint data-label log likelihoods (Grathwohl et al., 2020), where
z = logg((x y W. O)
(4)
where Z(W, 0) is normalizing constant, independent of the data. The term & E IR-Vv is zero-mean Gaussian noise term with variance Σ. Typically in logistic regression this noise term is ignored, although it has seen use model label noise (Collier et al., 2021). We include it to unify the presentation, and the variance can be assumed zero as necessary.
As in the regression case, we specify a Gaussian prior for W. In contrast with the regression setting, exact inference and computation of the posterior predictive is not analytically tractable in this model. We refer to this model- of multinominal Bayesian logistic regression on learned neural network features as discriminative classification, as logistic regression is a classical discriminative learning algorithm.
2.3 GENERATIVE CLASSIFICATION
The second classification model we consider is the generative classification model (Harrison et al., 2020; Zhang et al., 2021; Willes et al., 2022), so-called due to its similarity to classical generative models such as Gaussian discriminant analysis. In this model, we assume that the features associated with each class are normally distributed. Placing Normal prior the means of these feature distributions and (conjugate) Dirichlet prior on class probabilities, we have priors and likelihoods (top line and bottom line respectively) of the form
p Dir(a) p Cal(p)
By N Sy)
(5)
1 NN Σ).
(6)
In this model, Byy E No and S, NoxNN are the prior mean and covariance over By E R the mean embedding for each. The subscript here indexes the statistics for each class; we also write fr (, MN,) } to terms for all y. The terms p PN, correspond to class probabilities, where PN, denotes the probability simplex embedded in N,, These class probabilities are in turn used in the categorical distribution over the class.
We present scalar- output regression in the paper body and defer the multivariate output case to the appendix. 2 TTroughouu the paper, we use overbars denote mean parameters and underbars to denote prior parameters.
2

---

Published as a conference paper at ICLR 2024
15-
15
53-
03
0
3.5
9.5
-1.0
-1.
-15-
-7.5
-5.0
-2.5
0.0
2.5
5.0 7.5
Figure 1: Left: A variational BLL (VBLL) regression model with BBB features trained on 50 data points generated from cubic function with additive Gaussian noise. The plot shows the 95% predictive credible region under the variational posterior for several sampled feature weights. Right: Visualizing (re-scaled) p(æ y =) - p(æ y 0) predicted by generative VBLL model on the half moon dataset, shows good sensitivity to Euclidean distance and sensible embedding densities.
For distribution over model parameters
Ny p(p. μ n) = Dir(a) II ((xx,
for which write 77 {a, p.. S}. have
plu 1) = Cly
Σ p( y. 17) = NV D+ + Sy),
(8)
Σx= ax
via analytical marginalization. To compute the predictive over class labels, we apply Bayes' rule, yielding
p(y x, n) softmax, n)) + logp(y) n)).
Here,
log p(x y. m) - 1 ((0 ;y) (E Sy) -by) logdet(E + Sy) + +c)
(10)
where e constant, shared for all classes, that may be ignored due the shift-invariance of the softmax. Grouping the log determinant term with the class prior yields a bias term. Instead of a linear transformation of the input features to obtain class logit, we instead have quadratic transformation. This formulation is strict generalization of standard classifier architectures (Harrison, 2021), I), in which we have quadratic decision regions opposed to linear ones.
2.4 INFERENCE AND TRAINING IN BLL MODELS
BLL models have growing popularity in recent years, ironically driven in part by need for compatibility with increasingly deep models (Snoek al., 2015; Azizzadeneshelii al., 2018; Harrison et al., 2020; Ober Rasmussen, 2019; Kristiadi et al., 2020; Thakur et al., 2020; Watson et al., 2020; 2021; 1; Daxberger et al., 2021a; Willes et al., 2022; Sharma et al., 2022: Schwöbel et al., 2022; Zhang et al., 2021; Moberg et al., 2019; Fiedler & Lucia, 2023). Exact marginalization enables computationally efficient treatment of uncertainty, as well resulting in lower-variance training objectives compared to sampling-bee based Bayesian models. A common and principled objective for training BLL models the (log) marginal likelihood (Harrison et al., 2018), via gradient descent on
Harrison et al., 2018; Weber et al., 2018; Riquelme et al., 2018;
T-1 T-'log X,
(11)
where X, Y denote stacked data. We include a factor of 7 to enable better comparison with standard, non- non-Bayesian training pipelines (typically based on average loss over mini- mini-batches) and across dataset sizes. This training objective can be problematic, however: gradient computation requires computing the full marginal likelihood, and mini-batches do not yield unbiased gradient estimators as in standard training with an arbitrary loss function. Even mini-batch processing of the iterating between conditioning on mini-batches and prediction under the partial posterior induces long computation graphs that make training at scale impossible. Morcover, due to the flexibility of neural network features, full marginal likelihood training objective can result in substantial over-concentration of the approximate posterior (Thakur al., 2020; Ober et al., 2021).
3

---

Published as a conference paper at ICLR 2024
3 SAMPLING-FRES VARIATIONAL INFERENCE FOR BLL NETWORKS
To exploit exact marginalization while avoiding full marginal likelihood computation, we will turn to stochastic variational inference (Hoffman et al., 2013). In particular, we aim to jointly compute an approximate last layer posterior and optimize network weights by maximizing lower bounds on marginal likelihood. As such. we will avoid distributional assumptions made the previous section. We write the (uncertain) last layer parameters as & and aim to find an approximate posterior q(E 7) parameterized by 17. Concretely, throughout this section we will develop bounds of the form
T po X, 0) 2 [ n, Σ) T (KLEE n))lp({))
(12)
where L is architecture dependent and developed in the remainder of this section. Thus, practically, the T factor weights regularization terms in our training objective. In this section, we index data with (via subscript), including Φi := ( 0).
3.1 REGRESSION
We consider the log marginal likelihood log p(Y x, with marginalized parameters & = {w]. and have the following lower bound.
Theorem 1. Let q(E n) = Ni .) denote the variational posterior for the BLL. model defined in Section 1. Then, (12) holds with
T
(oggN(ye w Σ T
[(O, n, 2)
Φz 2) -
(13)
The proof for this result and all others is available in Appendix F. When g 1) P(E X) and distributional assumptions are satisfied, this lower bound is tight (this may be shown by direct substitution). This correspondence between the variational and true posterior for appropriately-chosen variational families is well known-sce Knoblauch al. (2019) for thorough discussion. We note that a similar objective for regression models was developed in Watson et al. (2021).
3.2 DISCRIMINATIVE CLASSIFICATION
In the discriminative classification case, the parameters are & = (W). We will assume diagonal covariance matrix Σ, and write of == Egg We will fix a variational posterior the form <(W n)) - Ny g(w; m) = -I qf (xx S&). where wy denotes the k'th row of W. This factorization retains dense covariances for each class, but sacrifices cross-class covariances. While only present this factorized variational posterior, similar training objective may be derived with fully dense variational posterior. Under the variational posterior, we have the following bound on the marginal likelihood.
Theorem 2.
Let q(W) | 1n) = =ITh ((wk, Sx) denote the variational posterior for the discrimina- live classification model defined in Section 2.2. Then, (12) holds with
1 T
1 y. Wo. LSE; wy Φ. + 2 (0, + + 1)
L(O, n. Σ)
Sgb.
(14)
Here, LSEx (-) denotes the log-sun sum--xxp function, with the sum over k. In contrast to the regression on the standard ELBO (due to two applications Jensen's 's inequality) and the bound is not tight. We have reduced variance (which would be induced by sampling logit values before the softmax in standard SVI (Ovadia et al., 2019)) for bias due to this lower bound. Our proof leverages the same double application of Jensen's inequality used by Blei & Lafferty (2007). We note that tighter analytically tractable lower bounds exist for the logistic regression model (Depraetere Vandebroek, 2017: 7; Knowles & Minka, 2011), although for simplicity of the resulting algorithm we use the above lower bound.
case, this lower bound is a lower
3.3 GENERATIVE CLASSIFICATION
In the generative classification case, the parameters are & [14, p}. ). In this setting, the Dirichlet posterior over class probabilities p(p Y) can be computed exactly with one pass over the data by simply counting class occurrences. We therefore only consider a variational posterior of the form
q(E n. Y) q(f=) 17) for the class embeddings, where q(1a n)) = II Mir,, Sx). This yields the following lower bound.
4

---

Published as a conference paper at ICLR 2024
Theorem 3. Let q/p n) - -1= -(p
N(ps, S&)) denote the variational posterior over class em-
beddings for the generative classification model defined in Section 2.3. Let p(p Y) = Dir(a) denote the exact Dirichlet posterior over class probabilities, with a denoling the Dirichlet posterior concentration parameters. Then, (12) holds with
C(O, n. -)
) + ((y))) Σ (log Py: 2) -01(= (yy))
(ax)) + log ax (15)
LLg Noo: p2 + Sk) + log oxs)
where (b)) is the digamma function and where ax, Σx OC
Importantly, we note that ((ay ) vb(a), ), log at, all vanish in gradient computation and may be ignored. The term log Org is the LSE can not be ignored, however. This training objective is again a lower bound on the ELBO, and is not tight. The first Dirichlet term (in the upper line) vanishes in gradient computation, but the second term inside the lug-sum-exp function does not. In the case that the posterior concentration parameters are equal for all classes (as in the case of balanced dataset), the concentration parameter can be pulled out of the LSB(.) (due to the equivariance of log-sum-exp under shifts) and can be ignored.
3.4 TRAINING VBLL MODELS
We propose three methods learn VBLL models.
Full training. First, we can jointly optimize the last layer variational posterior together with MAP estimation of the features, yielding combined training objective
0' ." E = arg max (Z(0, v), Σ) + (log + log p(E) - KL(((E 1) (P(())) (16) 0.
While one may expect this to result in substantial vvecconncnttatioon for weak feature priors, in practice we observe that stochastic regularization due to mini-batch optimization prevents overconcen- tration. Throughout this work, we will place simple isotropic zero-mean Gaussian priors on feature weights (yielding weight decay regularization) and a canonical inverse- Wishart prior on Σ. For Gaussian priors (as developed throughout this section) the KL. regularization term can be computed in closed form. The prior terms (and the KL penalty) introduce a set of new syperparameters that may be difficult to select. In Appendix C, we discuss these hyperparameters and their interpretation, and provide reformulation of hyperparameters that increases interpretability.
Post-training. As an alternative to jointly optimizing the variational last layer with the features, a two step procedure can used. In this step, the feature weights O are trained by any arbitrary training procedure (e.g. standard neural network training) and the last layer (and Σ) are trained with frozen features. The training objective is identical to (16), although 0 is trained in the initial pre-training step and n7' Σ' are trained via (16). Feature uncertainty. Lastly, can combine last layer SVI with variational feature learning (Blundell et al., 2015), corresponding to approximate collapsed VI (Teh et al., 2006). This training strategy allows us to construct a variational posterior on the full marginal likelihood, via
PoY X)= 2 (. Y E 0, Σ)] KL(g(E. 0, E 7) [P(E. O) (17)
Assuming the prior and variational posterior factorize across the features and last layer, we can partially collapse this expectation
Eg(g,0,r)m)) [Yo((y) X. 0. Σ)] = E,(00.7mm) [Yog(Y) X, 0, Σ)] 2 TE (C(E, E)
(18)
and the KL penalty may be similarly decomposed into several terms that be computed in closed form under straightforward distributional assumptions. In the above, we have included Σ in the variational posterior, although practically we perform MAP estimation of this covariance under inverse- priors. Again in this setting, pre-training and post-training steps may be combined, but we do not investigate this case.
3.5 PREDICTION WITH VBLI MODELS
For prediction in VBLL models, we will predict under the variational posterior directly, approximating (for test input/label (x, z)).
(19) for the deterministic feature model. This expectation may be computed in closed form (for the regression and generative classification model) due to conjugacy, and can be computed via inexpensive last layer sampling in the discriminative classification model. In the variational feature model,
p(y x, X,Y) R (p( x, &, 0" E"'
p(y x. X, Y) Ey(( Eygelm (ply x. 0, Σ')
(20)
5

---

Published as a conference paper at ICLR 2024
where the inner expectation may be computed exactly and the outer expectation may be approximated via sampling. Further details training, prediction, and out of distribution detection within all three VBLL models is provided in Appendix B.
For both training and prediction, under relatively weak assumptions on covariance matrices, com- putational complexity (for the classification models) is at most ONN, and can be reduced to ONN, ) for diagonal covariances. This matches the complexity of standard network evaluation; for reasonable choices of covariance sparsity, the additional computational cost of VBLL models over standard networks is negligible. More details are provided in Appendix C.
4 RELATED WORK AND DISCUSSION
Bayesian methods capable of flexible nonlinear learning have been a topic of active study for the last several decades. Historically, early interest in Bayesian neural networks (MacKay, 1992; Neal, 1995) diminished as Gaussian processes rose to prominence (Rasmussen, 2004). In recent years, however, there has been growing interest in methods capable of learning expressive features, effectively quantifying uncertainty, and training efficiently on large datasets. Variational methods have seen particular attention in both neural networks (Blundell et al., 2015; Ovadia et al., 2019) and GPs (Hensman et al., 2013; Titsias, 2009; Liu et al., 2020) due to their flexibility and their ability to produce mini-batch gradient estimation training schemes.
While a wide range of work has aimed to produce more performant approximate Bayesian methods (including more expressive prior and posterior representations (Fortuin et al., 2021; Izmailov et al., 2021; Sun et al., 2019; Wilson & Izmailov, 2020)), they have still seen limited application, often due to the increased computational expense of these methods (Lakshminarayanan et al., 2017: Dusenberry et al., 2020). While some approaches to Bayesian neural networks have focused on improving the quality of the posterior uncertainty through e.g. better priors (Farquhar et al., 2020; Fortuin, 2022) or inference methods (Izmailov et al., 2021). other lines of work have focused on designing comparatively inexpensive upproximate Bayesian methods. Indeed, simple strategies such a
as Bayesian dropout (Gal& & Ghahramani, 2016) and stochastic weight averaging (Maddox et al., 2019) have seen much wider use than more expressive methods due to their simplicity.
One of the simplest Bayesian models is the BLL model that is the focus of this paper, which enables single-pass, often deterministic uncertainty prediction. This model has gained prominence through the lens of deep kernel learning (Wilson al., 2016b;a; Watson et al., 2020; Liu et al., 2022) and within few-shot learning (Harrison et al., 2018: 2020; Harrison, 2021; Watson et al., 2021; Zhang et al., 2021). Deep kernel learning aims to augment standard neural network kernels with neural network inputs. This approach allows control of the behavior of uncertainty, particularly as a function of Euclidean distance (Liu et al., 2022). While stochastic variational inference has been applied to these models (Wilson et al., 2016a), efficient and deterministic mini-batch methods have not been major focus. Moreover, classification softmax functions, which increases variance (Ovadia et al., 2019; Kristiadi et al., 2020; 202 1), or on Laplace approximation (Liu et al., 2022).
these models typically relies on sampling logits applying
Within few-shot learning, exact conjugacy of the Bayesian linear regression model (Harrison et al., 2018) and Bayesian GDA (Harrison et al., 2020; Zhang et al., 2021; Snell et al., 2017) has been exploited for efficient few-shot adaptation. These models have (in addition to Van Amersfoort et al. (2020) among others) shown the strong performance of baseebase/raaial basis function networks, especially on problems such as out of distribution detection, which we further highlight in this work. However, training these models (as well as the DKL methods discussed previously) relies on direct computation of the marginal likelihood. In contrast to prior work on DKL and few-shot learning, our approach achieves efficient and deterministic training and prediction through our variational objectives and through similarly exploiting conjugacy, and thus the added complexity compared to standard neural network models minimal.
5 EXPERIMENTS
We investigate the three VBLL models, with both MAP and variational feature learning, in regression and classification tasks. full description of all metries used throughout this section and baseline methods is available in the appendix. To illustrate VBLL models, we show predictions on simple datasets in Figure 1. The left figure shows a regression VBLL model with variational features trained on the function f(x) = cx², with training data shown in red. This figure shows the behavior on
Complexity for the regression case is ONN3 + N2).

---

Published as a conference paper at ICLR 2024
Table I: Results for UCI regression tasks.
BOSTON
CONCRETE
ENERGY
NL).(L)
RMSE ())
N.)
RMSE ())
NIL
RMSE (L)
VBLL
0.06
2.92 0.12
3.22 0.07
5.09 0.13
1.37 - 0.08
0.87
GBLL
2.90
4.19
3.09
0.03
5.01 + 0.18
0.69 - 0.03
0
LDGBLL
2.00
3.38
2.97 0.03
4.80
4.18
0.02
MAP
0.07
0.11
3.04 0.04
0.12
0.04
0.00
RBF GP
2.41
3.08 0.16
5.62 0.13
0.64
0.00
Dropout
2.36
0.04
0.16
2.90 0.02
4.45
1.30
0.53 0.001
Ensemble
2.48 0.09
0.17
0.04
0.12
0.58
0.41
SWAG
2.64 0.16
3.08
0.35
3.19
5.10
0.08
0.09
BBB
2.39 0.04
0.16
2.97 + 0.03
4.80
0.63
0.01
VBLL BBB
2.59 + 0.07
3.13 11130.10 0.19
3.36 + 0.22
5.16
1.35 0.15
0.062 0.03
Table 2: Further results for UCI regression tasks.
POWER
WINE
YACHT
NL)
RMSE (L)
NLL (L)
RMSE (J)
NLL (J)
RMSE (L)
VBLL
3.73#0001
3.68 - 0.03
1,02 1 0.03
0.65 + 0.01
1.19
0.86 0.86::11 0.17
GBLL
2.77
1.02 + 0.001
0.64 + 0.01
1.67 0.11
1.09 0.09
LDGBLL
2.77
0.04
1.02 0.01
0.64+ 0.01
+.13+0.06
0.75
MAP
0.00
3.01
0.96 0.01
5.14
0.91
RBF GP
3.72+
0.45 0.01
0.56 0.05
217+0.00
0.40 0.03
Dropout
2.80
0.93 0.01
+.00
1.02
1.21 +.13
Ensemble
+.00
3.59
0.95 0.01
0.01
0.07
0.08
SWAG
2.77
3.85 + 0.05
0.96 + 0.03
0.63+
1.11 0.05
1.13+00.21
BBB
2.77 0.01
3.86 0.04
0.95 + 0.01
0.63 + 0.01
1.43 0.17
1.10
VBLL BBB
2.74 0.01
3.73+0.00
0.94 0.03
0.61 + 0.01
2.96
+ 0.59
0.79
so-called gap datasets- SO named because of the interval between subsets of the data. The VBLL model shows desirable increasing uncertainty between the intervals (Foong et al., 2019). The right figure shows the generative classification model (G-VBLL) on the half-moon dataset. In particular, we visualize the feature density for each class. Importantly, the density has high Euclidean distance sensitivity, which has been advocated by Liu et al. (2022) as a desirable feature for robustness and out of distribution detection.
5.1 REGRESSION
We investigate the performance of the regression VBLL models on UCI regression datasets (Dua & Graff, 2017), which are standard benchmarks for Bayesian neural network regression (Moberg et al., 2019; Ober & Rasmussen, 2019; Daxberger et al., 202 Watson et al., 2021; Kristiadi et al., 2021). Results are shown in Tables 1, 2. We include baseline models run in Watson et al. (2021), and we replicate their experimental procedure and hyperparameters as closely as possible (details in the appendix).
Our experiments show strong results for VBLL models across datasets. Of particular interest is the performance relative to the GBLL model, which is trained directly on the exact marginal likelihood within the Bayesian last layer model. There are several contributing factors: the prior parameters were jointly optimized with the feature weights in the GBLL model, whereas prior terms were fixed in our VBI.I. model, resulting inference can perform poorly under model misspecification (Grünwald & Van Ommen, 2017), whereas variational Bayes has comparatively favorable robustness properties and asymptotics (Giordano et al., 2018; Wang & Blei, 2019), although the Gaussian process (GP) model generally also has strong performance across datasets. Finally, directly targeting the marginal likelihood (computed exactly within conjugate models such as BLL models) has been shown to induce substantial overfitting (Ober et al., 2021, ;; Thakur et al., 2020; Harrison, 2021), which the variational approach may avoid due to worse inferential efficiency.
stronger regularization effect. Moreover, exact Bayesian
5.2 IMAGE CLASSIFICATION
To evaluate performance of VBLL models in classification, we train the discriminative (D-VBLL) and generative (G-LBLL) classification models on the CIFAR- 10 and CIFAR-100 image classification task. Following Liu et al. (2022), all experiments utilize a Wide ResNet-28-11 backbone architecture. We investigate full training methods (without post-training step), indicated with the method name in the top third of Tables 3, 4: post-training methods, indicated by pre-training method + post-training method, in the middle third of the Tables; and feature uncertainty, in the bottom third.
We evaluate out of distribution (OOD) detection performance using the Street View House Numbers (SVHN) (Netzer et al., 2011) as a far-OOD dataset for both datasets, and CIFAR- 100 for CIFAR- 10 (and vice-versa) as near-OOD datasets. In-distribution data normalization is used in both cases. The DNN, BBB, D-VBLL and D--BBL BBB models maximum softmax probability (Hendrycks &
7

---

Published as a conference paper at ICLR 2024
Table 3: Results for Wide ResNett-2-100 on CIFAR- 10.
Method
Accuracy (1)
FCR (L)
SVHN AUC (
CIFAR-100 AUC (1)
DNN
95.8
10028:0:025
0.007
0.936 0.005
0.893 0.001
SNGP
0.17
0.017 0.003
4.149:0.000
0.902 - 0.003
D-VBLL
9
0.022
0.001
0.160
0.969 0.004
0.900 + 0.004
G-VBLL
96.3 0.06
0.021 0.000
0.002
0.005
0.804 0.006
DNN + LL Laplace
0.03
0.000
0.133+0..00
0.963
U.898 = 0,001
DNN + D-VBLL
96.4
0.024 0.000
0.000
0.002
0.895 + 0.000
DNN + G-VBLL
9..4
0.035 + 0.000
0.533 + 0.003
0.729
0.661 + 0.004
G-VBLL + MAP
100006
0.893 1 0.003
Dropout
0.17+
0.013 0.002
0.145 + 0.004
0.934
0.903 0.001
Ensemble BEB
9.04
0.011 0.092
0.124 + 0.001
0.947 0.002
0.914 - 0.000
9.00
0.033 + 0.001
0.333 0.014
0.957 + 0.004
0.844 + 0.013
D-VBLL BBB
95.9
0.058 1 0.019
0.238 1 0.036
0.832 1 0.026
0.744 1 0.010
G-VBLL BBB
95.9 0.166
0.009 1 0.001
0.229 1 0.010
0.017 1 0.005
0.779 + 0.009
Table 4: Results for Wide ResNet-28-10 on CIFAR-100,
Method
Accuracy (1)
ECE (1)
N.L
SVHN AUC (1)
CIFAR-IO AUC (1)
DNN
80.4 0.29
0.107 0.004
0.799 = 0.020
0.795 = 0.001
SNGP
80.3 + 0.23
0.030 0.004
1761:0000
0.846 0.019
0.798 + 0.001
D-VBLL
80.7 0.03
0.000
0.913
0,849 0.006
0.791 = 0.003
G-VBLI.
80.4 0.10
0.051 0.003
0.945+
0.767 0.055
0.752 0.015
DNN + LL Laplace
80.4 0.29
0.210
0 0.01
0.811 = 0.002
DNN + D-VBLL
80.7 + 0.02
0.063 1: 0.000
1.31
1.000100001
0.804 1 0.001
DNN + G-VBLL
80.6 0.02
0.186 0.003
3.026 0.155
0.638 0.021
0.652 0.025
G-VBLL + MAP
0.793 + 0.032
0.765 + 0.008
Dropout
0.22+0.22
0.031 + 0.002
0.762 + 0.008
0.800 1:000+000 0.014
0.797 + 0.002
Ensemble
82.5 +.5=00 0.19
0.041 T 0.002
0.674 + 0.004
0.812 1 0.007
0.814 - 0.001
BBB
79.6 - 0.04
0.127 + 0.002
1.611 + 0.006
0.809 + 0.060
0.777 + U.00S
D-VBLL BBB
77.6 0.17
0.041 0.003
1.169 0.018
0.785 0.022
0.756 0.002
G-VBLL BBB
78.1 + 0.18
0,046 1 0.002
1.156 1 0.008
0.832 + 0.023
0.742 1 0.004
Gimpel, 2016) as an OOD measure. The G-VBLL and G-VB BBB models use normalized feature density. Two methods for this exist: G-VBLL and G-VBLL BBB both use the learned variational posteriors to compute feature likelihoods. However, the performance of this is relatively weak, as there is no guarantee that learned feature likelihoods correspond effectively to true embedding densities. Thus, we also investigate an approach which we estimate distributions for fixed features after training. This method estimates noise covariances for each class using the trained features, similar to the approach used in Liu et al. (2022). We refer to this model as G-VBLL-MAP as the approach corresponds to MAP noise covariance estimation. These estimated covariances often result in overly overly-confident predictions, and so we do not advocate for label prediction under these fit covariances, and do not include results for them. Appendix B.6 discusses OOD methods, and further experimental details are in Appendix D.
Tables 3,4 summarize the CI and CIFAR- 100 results. D-VBLL and G-VBLL report strong accuracy performance and competitive metrics for both ECE and NLL. D-VBLL in particular demonstrates strong accuracy results, as well as competitive (with SNGP) NLL and OOD detection ability. Despite its comparative simplicity. it outperforms SNGP on accuracy and OOD on CIFAR-10 and accuracy on CIFAR- -00. It matches SNGP on OOD for CIFAR-100 and is competitive (although slightly worse) on ECE and NLL. Overall, D-VBLL models stand out for their strong performance relative to their complexity. They also perform well as post-training models, whereas G-VBLL performs is substantially degraded.
While models with MAP feature estimation show strong performance versus baseline models, the performance of variational feature learning models (BBB) is more mixed. In regression tasks, these models are competitive, while classification the performance is worse than deterministic models. In both settings, we use default KL term weighting (one over the dataset size). This contrasts with the tempered/cold posterior effect (Kapoor et al., 2022; Wenzel et al., 2020; Izmailov et al., 2021; Aitchison, 2020), in which it has been observed that alternative weightings of the likelihood and the KL may outperform this one. This is attributable (in part) to two factors: data augmentation and stochastic regularization. In regression there is no data augmentation and the model is trained for substantially longer than deterministic models; in classification we use standard augmentation and our training is more limited. Thus, it is possible that classification BBB models are weer-regulaiize.. We investigate this question in more detail in the appendix.
5.3 SENTIMENT CLASSIFICATION WITH LLM FEATURES
We evaluate VBLL models for language modelling tasks using the IMDB Sentiment Classification Dataset (Maas et al., 2011) 1). The IMBB dataset is a binary text classification task consisting of 25,000 polarized movie reviews for training and another 25,00 25,000 for testing. A pre-trained OPT-175B (Zhang et al., 2022) model is used for text feature extraction. Sequence embeddings are obtained from OPT
as the last token output from the the final network layer. We train both the generative (G-VBLL) and
8

---

Published as a conference paper at ICLR 2024
0.30
MUP DAYBLL DIBBLL
90
0.25
Aruney
Fredetive
0.15
0.6
60
0.4
0.05
0.2
0.00
10
Figure 2: A performance comparison of G-BLL, D-VBLI and baseline MLP models on the IMDB Sentiment Classification Dataset. The models utilize text embeddings extracted from pre-trained OPT-175B model. Results are presented across multiple training dataset scales, and the shaded regions represent la error bounds.
Table 5: Wheel bandit cumulative regret.
8
8 = 0.7
6 0.9
6= 0.95
8 0.99
VALL
0.46 105::0.00
- 0.01
2.54
0.02
0.03
0.71
NeuralLincar
1.10 0.00
1.77 4 0.03
1.32 = 0.11
11.42
+ 0.97
52.64 + 2.01
Neurall Lineur MR
0.95 +
1.60
1 0.03
4.65
+ 0.18
9.56
0.36
49.63 2.41
PosDiag
1.12 11220.0 0.03
1.80 + 0.08
5.06 +
8.90
37.77
Table 6: Wheel bandit simple regret.
5 = 0.5
5 0.7
8= 0.9
5 0.95
8 = 0.99
VRI.I
0.27 4 0.03
0.69
0.06
2.28
4.16
0.17
21.05 + 1.59
NeuralLinear
0.31 - 0.03
0.68 = 0.07
2.18 0.13
5.44
0.73
46.42 3.45
NeuralLimear- MR
0.33 - 0.04
0.79 1 0.07
2.17 0.14
4.08
+ 0.20
35.89 1: 2.98
LinPost MR
0.70 0.06
0.99
1 0.10
3.08 0.22
4.85
0.27
25.42 1.81
discriminative (D-VBLL) models and a baseline MLP on the sequence embeddings via supervised learning at multiple training dataset scales: 10, 100, 1000 and 25,000 training samples. Evaluation is performed using the complete test set at each training dataset scale. Results are shown in Figure 2. The VBLL models demonstrate strong performance significantly lower predictive NLL and ECE smaller training dataset sizes. These findings validate the VBLL models' potential for integration with large-scale modern language models for diverse applications, particularly in sentiment classification tasks.
comparison the MLP baseline. We see
5.4 WHEEL BANDIT
To investigate the value of VBI.I. models in an active learning setting, we apply VBLI. regression model to the wheel bandit problem presented in Riquelme et al. (2018). This problem is contextual bandit in which the state is sampled randomly in two dimensional ball, and the learned model aims to identify the reward function. There are five regions in the ball and five actions: each region roughly corresponds to a correct action yielding a high reward, and incorrect action choice yields low reward, although action 1 always yields an intermediate reward and no high-reward action exists for region 1. The parameter ó controls the volume of the high-reward regions, with larger d corresponding smaller high-reward regions. We report both cumulative regret- the difference in reward compared to an oracle, normalized to the performance of a random agent, aggregated over the full problem duration and the simple regret, which captures only the last 500 timesteps and thus (roughly) measures the final quality of the learned model. We use a Thompson sampling policy (Russo et al., 2018; Thompson, 1933), and compare to the top models reported in (Riquelme et al., 20. We find that our VBLL model strongly outperforms the top performing baselines in cumulative regret (Table 5) and slightly outperforms them in simple regret (Table 6), implying both the capacity of the model matches the best baselines while also exploring more effectively.
CONCLUSIONS AND FUTURE WORK
We have presented a simple, nearly computationally free Bayesian last layer architecture that can be applied to arbitrary network backbones. The practical realization of the VBLL model is a small number of extra parameters (corresponding to the variational posterior covariance) and a small number of regularization terms corresponding to terms arising in the marginalized predictive likelihood, prior terms used in MAP estimation, and KL divergences. Several important directions for future work exist. First, few-show adaptation that further exploits conjugacy of these models via e.g. recursive Bayesian least squares is possible. We have only leveraged basic ideas from variational inference in this work: there are many highly practical ideas within variational Kalman filtering which may enable efficient model adaptation, label noise robustness, inference within heavy-tailed noise, or improved time series filtering (Sykacek & Roberts, 2002; Sarkka & Nummenmaa, 2009; Ting et al., 2007).

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
We acknowledge Apoorva Sharma, Jascha Sohl-Dickstein, Alex Alemi, and Allan Zhou for useful conversations over the course of this work. We also gratefully acknowledge Paul Brunzema, who identified a subtle bug in our initial results.
REFERENCES
Laurence Aitchison. A statistical theory of cold posteriors in deep neural networks. arXiv preprint wrXiv:2008.05912 2020.
Kamyar Arizzadenesheli, Emma Brunskill, and Animashree Anandkumar. Efficient exploration through Bayesian deep q-networks. 180. 2018.
David Blackwell. Conditional expectation and unbiased sequential estimation. The Annals of Mathematical Statistics, 1947.
David M Blei and John D Lafferty. A correlated topic model of science. The annals of applied statistics, 2007.
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International Conference Machine Leurning (ICML), 2015.
George EP Box and George Tiao. Bayesian inference in statistical analysis, volume 40. John
Wiley & Sons, 2011.
Mark Collier, Basil Mustafa, Efi Kokiopoulou, Rodolphe Jenatton, and Jesse Berent. Correlated input-dependent label noise in large-scale image classification.. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2021.
Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace redux-effortless. Bayesian deep learning. Neural Information Processing Systems (NeurIPS), 2021a.
Erik Daxberger, Eric Nalisnick, James U Allingham, Javier Antorán, and José Miguel Hernández- Lobato. Bayesian deep learning via subnetwork inference. In International Conference on Machine Learning (ICML), pp. 25 1010215... PMLR, 2021b.
Nicolas Depraetere and Martina Vandebroek. A comparison of variational approximations for fast inference in mixed logit models. Computational Statistics, 2017.
Dheeru Dua and Casey Graff. UCI machine learning repository, 2017. URL http::/arcchive.. ics. uci. www.eeu..meeee eedu/m..
Michael Dusenberry, Ghassen Jerfel, Yeming Wen, Yian Ma, Jasper Snock, Katherine Heller, Balaji Lakshminarayanan, and Dustin Tran. Efficient and scalable bayesian neural nets with rank-1 factors. In International Conference on Machine Learning (ICML), 2020.
Sebastian Farquhar, Michael A Osborne, and Yarin Gal. Radial bayesian neural networks: Beyond discrete support in large-scale bayesian deep learning. In Artificial Intelligence and Statistics (AISTATS), 2020.
Felix Fiedler and Sergio Lucia. Improved uncertainty quantification for neural networks with bayesian last layer. arXiv preprint arXiv:2302. 2023.
Andrew YK Foong, Yingzhen Li, José Miguel Germández-Lobate and Richard E Turner. 'in-between' uncertainty in bayesian neural networks. arXiv preprint arXiv: 1906. 1537, 2019.
Vincent Fortuin. Priors in bayesian deep learning: A review. International Statistical Review, 2022.
Vincent Fortuin, Adrià Garriga-Alonss Sebastian W Ober, Florian Wenzel, Gunnar Rätsch, Richard Turner, Mark van der Wilk, and Laurence Aitchison. Bayesian neural network priors revisited. arXiv:2102.0657/, 2021.
Yarin Gal and Zoubin Ghahramani. Dropout as bayesian approximation: Representing model uncertainty in deep learning. In International Conference on Machine Learning (ICML), 2016.
Seymour Geisser. Bayesian estimation in multivariate analysis. The Annals of Mathematical Statistics, 1965.
10

---

Published as a conference paper at ICLR 2024
Ryan Giordano, Tamara Broderick, and Michael I Jordan. Covariances, robustness and variational bayes. Journal of Machine Learning Research, 2018.
Will Grathwohl, Kuan-Chich Wang.
Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky. Your classifier is secretly an energy based model and you should treat it like one. International Conference on Learning Representations (ICLR), 2020.
Peter Grünwald and Thijs Van Ommen. Inconsistency of bayesian inference for misspecified linear models, and a proposal for repairing it. Bayesian Analysis, 2017.
James Harrison. Uncertainty and efficiency in adaptive robot learning and control. PhD thesis, Stanford University, 2021.
James Harrison, Apoorva Sharma, and Marco Pavone. Meta-learning priors for efficient online Bayesian regression. Workshop on the Algorithmic Foundations of Robotics (WAFR), 2018.
James Harrison, Apoorva Sharma, Chelsea Finn, and Marco Pavone. Continuous meta-learning without tasks. Neural Information Processing Systems (NeurIPS), 2020.
Dan Hendrycks and Kevin Gimpel. A baseline for detecting misclassified and out-of-distribution examples in neural networks. arXiv:1610.02136. 2016.
James Hensman, Nicolo Fusi, and Neil D Lawrence. Gaussian processes for big data. Uncertainty in Artificial Intelligence (UAT), 2013.
Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley. Stochastic variational inference. Journal of Machine Learning Research, 2013.
Pavel Izmailov, Sharad Vikram, Matthew D Hoffman, and Andrew Gordon Gordon Wilson. What are bayesian neural network posteriors really like? In International Conference on Machine Learning (ICML), 2021.
David Jacobson.
Optimal stochastic linear systems with exponential performance criteria and their relation deterministic differential games. IEEE Transactions Automatic Control, 1973.
Sanyam Kapoor, Wesley Maddox. Pavel Izmailov, and Andrew G Wilson. On uncertainty, tempering. and data augmentation in bayesian classification. In Neural Information Processing Systems (NeurIPS), 2022.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations (ICLR), 2015.
Diederik Kingma and Max Welling, Auto-encoding variational Bayes. International Conference on Learning Representations (ICLR), 2014.
Jeremias Knoblauch, Jack Jewson, and Theodoros Damoulas. Generalized variational inference:
Three arguments for deriving new postcriors. arXiv preprint arXiv: 2019.
David Knowles and Tom Minka. Non-conjugate variational message passing for multinomial and binary regression. In Neural Information Processing Systems (NeurIPS), 2011.
Agustinus Kristiadi, Matthias Hein, and Philipp Hennig. Being Bayesian, even just a bit, fixes overconfidence in relu networks. In International Conference on Machine Learning (ICML), 2020.
Agustinus Kristiadi, Matthias Hein, and Philipp Hennig. Learnable uncertainty under laplace approximations, In Uncertainty in Artificial Intelligence (UAI), 2021.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Neural Information Processing Systems (NeurIPS), 2017.
Haitao Liu, Yew-Soon Ong, Xiaobo Shen, and Jianfei Cai. When gaussian process meets big data: A review of scalable gps. IEEE transactions neural networks and learning systems, 2020.
Jeremiah Liu, Shreyas Padhy, Jie Ren, Zi Lin, Yeming Wen, Ghassen Jerfel, Zack Nado, Jasper Snoek, Dustin Tran, and Balaji Lakshminarayanan. A simple approach to improve single-model deep uncertainty via distance-awarenes. Journal of Machine Learning Research, 2022.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization.
2017.
11

---

Published as a conference paper at ICLR 2024
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng. and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguisties: Human Language Technologies, pp. 142-150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics.
David JC MacKay, A practical bayesian framework for backpropagation networks. Neural Computa- lion, 1992.
Westey Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson A simple baseline for bayesian uncertainty in deep learning. In Neural Information Processing Systems (NeurIPS), volume 32, 2019.
John Moberg, Lennart Svensson, Juliano Pinto, and Henk Wymeersch. Bayesian linear regression on 2019.
deep representations. arXiv:/9/2.061
Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih. Monte carlo gradient estimation in machine learning. Journal of Machine Learning Research, 2020.
Eric Thomas Nalisnick. On priors for Buyesian neural networks. University of California, Irvine, 2018.
Radford M Neal. Bayesian Learning for Neural Networks. PhD thesis, University of Toronto, 1995.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y. Ng. Reading digits in natural images with unsupervised feature learning, In NIPS Workshop on Deep Learning and Unsupervised Feature Learning 2011, 2011. URL http:/ housenumbers / ni mipps0111111 www.hhuuuuuuuss...o.o pdf.
http:://fldd.... /ufldll stanford. edu/
Sebastian W Ober and Carl Edward Rasmussen. Benchmarking the neural linear model for regression. arXiv preprint arXiv: arXiv:1912.088 1912.0881,, 2019.
Sebastian W Ober, Carl E Rasmussen, and Mark van der Wilk. The promises and pitfalls of deep kernel learning. In Uncertainty in Artificial Intelligence (UAI), 2021.
Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, Sebastian Nowozin, Joshua Dillon, Balaji Lakshminarayanan, and Jasper Snoek. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. In Neural Information Processing Systems (NeurIPS), 2019.
Theodore Papamarkou, Jacob Hinkle, M Todd Young, and David Womble, Challenges in markov chain monte carlo for bayesian neural networks. Statistical Science, 2022.
Tim Pearce, Russell Tsuchida, Mohamed Zaki, Alexandra Brintrup, and Andy Neely. Expressive priors in Bayesian neural networks: Kernel combinations and periodic functions. In Uncertainty in Artificial Intelligence (UAI), 2020.
Fabian Pedregosa, Gaël Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, et al. Scikit-learn: Machine learning in python. the Journal of machine Learning research, 2011.
C Radhakrishna Rao. Information and the accuracy attainable in the estimation of statistical parame- ters. In Breakthroughs in statistics. Springer, 1992.
Carl Edward Rasmussen. Gaussian processes in machine learning. Springer, 2004.
Jie Ren, Peter Liu, Emily Fertig, Jasper Snoek, Ryan Poplin, Mark Depristo, Joshua Dillon, and Balaji Lakshminarayanan. Likelihood ratios for out-of-distribution detection. In Neural Information Processing Systems (NeurIPS), 2019.
Jie Ren, Stanislav Fort, Jeremiah Liu, Abhijit Guha Roy, Shreyas Padhy, and Balaji Lakshminnrayaaann A simple fix to mahalanobis distance for improving near-ood detection. arXiv:2106.09022, 2021.
Carlos Riquelme, George Tucker, and Jasper Snoek. Deep Bayesian bandits showdown. In Interna- tional Conference on Learning Representations (ICLR), 2018
Daniel Russo, Benjamin Van Roy, Abbas Kazerouni, Ian Osband, Zheng Wen, et al. A tutorial on Thompson sampling. Foundations and Trends in Machine Learning, 2018
Simo Sarkka and Aapo Nummenmaa. Recursive noise adaptive kalman filtering by variational bayesian approximations. IEEE Transactions Automatic Control, 2009.
12

---

Published as a conference paper at ICLR 2024
Pola Schwöbel, Martin Jørgensen, Sebastian W Ober, and Mark Van Der Wilk. Last layer marginal likelihood for invariance learning. In Artificial Intelligence and Statistics (AISTATS), 2022.
Mrinank Sharma, Sebastian Farquhar, Eric Nalisnick, and Tom Rainforth. Do Bayesian neural networks need to be fully stochastic? arXiv preprint arXiv:221 2022.
Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning. Neural Information Processing Systems (NeurIPS), 2017.
Jasper Snock, Oren Rippel. Kevin Swersky, Ryan Kiros, Nadathur Satish, Narayanan Sundaram, Mostofa Patwary, Prabhat, and Ryan Adams. Scalable Bayesian optimization using deep neural networks. International Conference on Machine Learning (ICMI), 2015.
Shengyang Sun, Guodong Zhang, Jiaxin Shi, and Roger Grosse. Functional variational bayesian neural networks. A0Kiv:1903.05.0. 2019.
Peter Sykacek and Stephen J Roberts. Adaptive classification by variational kalman filtering. In Neural Information Processing Systems (NeurIPS), 2002.
Yee Teh, David Newman, and Max Welling. A collapsed variational bayesian inference algorithm for latent dirichlet allocation. In Neural Information Processing Systems (NeurlPS), 2006.
Sujay Thakur, Cooper Lorsung, Yaniv Yacoby, Finale Doshi-Velez, and Weiwei Pan. Uncertainty- aware (una) bases for Bayesian regression using multi-headed auxiliary networks. arXiv preprint arXiv:2006. / 1695, 2020.
William R Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika, 1933.
George C Tiao and Amold Zellner. On the bayesian estimation of multivariate regression. Journal of the Royal Statistical Society; Series B (Methodological), 1964.
Jo-Anne Ting, Evangelos Theodorou, and Stefan Schaal. A kalman filter for robust outlier detection. In IEEE International Conference on Intelligent Robots and Systems (IROS), 2007.
Michalis Titsias. Variational learning of inducing variables in sparse gaussian processes. In Artificial Intelligence and Statisties (AISTATS), 2009.
Joost Van Amersfoort, Lewis Smith, Yee Whye Teh, and Yarin Gal. Uncertainty estimation using single deep deterministic neural network. In International Conference on Machine Learning (ICML., 2020.
Yixin Wang and David Blei. Variational bayes under model misspecification. Neural Information Processing Systems (NeurIPS), 2019.
Joe Watson, Jihao Andreas Lin, Pascal Klink, and Jan Peters. Neural linear models with functional gaussian process priors. In Advances in Approximate Bayesian Inference (AABI), 2020.
Joe Watson, Jihao Andreas Lin, Pascal Klink, Joni Pajarinen, and Jan Peters. Latent derivative Bayesian last layer networks. In Artificial Intelligence and Statistics (AISTATS), 2021.
Noah Weber, Janez Stare, Arpit Mittal, Roi Blanco, and Lluís Märquez. Optimizing over Bayesian last layer. In NeurlPS workshop on Bayesian Deep Learning, 2018
Florian Wenzel, Kevin Roth, Bastiaan S 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 international Conference on Machine Learning (ICML), 2020.
John Willes, James Harrison, Ali Harakeh, Chelsea Finn, Marco Pavone, and Steven Waslander. Bayesian embeddings for few-shot open world recognition. IEEE Transactions on Pattern Analysis & Machine Intelligence, 2022.
Andrew G Wilson and Pavel Izmailov. Bayesian deep learning and a probabilistic perspective of generalization. In Neural Information Processing Systems (NeurIPS), volume 33, 2020.
Andrew G Wilson, Zhiting Hu, Russ R Salakhutdinov, and Eric P Xing. Stochastic variational deep kernel learning. In Neural Information Processing Systems (NeurIPS), 2016a.
13

---

Published as a conference paper at ICLR 2024
Andrew Gordon Wilson, Zhiting Hu, Ruslan Salakhutdinov, and Eric P Xing. Deep kernel learning. In Artificial Intelligence and Statistics (AISTATS), 2016b.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. arXiv preprint arXvv:600..07146 2016.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trale ire-trained transformer language models. arXiv preprint arXiv:2205.0100 01068, 2022.
Xueting Zhang, Debin Meng, Henry Gouk, and Timothy M Hospedales. Shallow Bayesian meta learning for real-world few-shot recognition, In International Conference on Computer Vision (ECCV), 2021.
14

---

Published as a conference paper at ICLR 2024
A The Multivariate Regression Model
16
Algorithmic Details
16
B.1
Feature Point Estimation
16
B.2
Post-Training with VBLL Layers
16
B.3
Collapsed Variational Inference for Bayesian Neural Networks Training
16
B.4
17
B.5 Prediction and Monitoring
18
B.6 Out of Distribution Detection
19
Parameterization, Complexity, Regularization, and Hyperparameters
19
C.I
Regression Complexity
20
C.2 Classification Complexity
C33 Complexity of Comparable Baselines
21
C.4 Hyperparameters
21
C,5 Understanding Prior Regularizers
22
Experimental Details
23
D.I Metrics
23
D.2 Baselines
23
D.3 Toy Experiments
24
D.4 Regression
24
D.5 Image Classification
25
D.6 Sentiment Classification with LLM Features
25
D.7 Wheel Bandit
25
Hyperparameter Studies and Ablations
26
F Proofs and Further Theoretical Results
27
F.1 Helper Results
27
F2 Proof of Theorem
28
F3 Proof of Theorem 2
29
F.4 Proof of Theorem 3
29
15

---

Published as a conference paper at ICLR 2024
A THE MULTIVARIATE REGRESSION MODEL
In the multivariate regression case, we consider a model of the form
= Φ +
(21)
and place a matrix normal (Tiao & Zellner, 1964; Geisser, 1965) prior on W. with W MN (w. I, S). For a discussion of the matrix normal distribution, we refer the reader to (Box & Tiao, 2011).
Given the matrix normal prior and the above model, the posterior is also matrix normal. We thus fix matrix normal variational posterior. Appendix F.2, obtain ELBO of the form
T
1
Σ ==1
log (w W6 Σ) 20 (
L( Σ)
(22)
for 17 = ( 5}. and we use this as training objective.
For parameter distribution MW(WW, L, S), prediction in this model is analytically tractable and is p(y: xx, 77, 0) NNW 0, S L).
(23)
B ALGORITHMIC DETAILS
In this section we present concrete details on training VBLL models. We first describe the procedure for MAP estimation, last layer training on frozen features, and variational learning of features, as described in the paper body. We then discuss prior choice, describe the resultant regularization terms, and describe prediction and out of distribution detection within these models.
B. FEATURE POINT ESTIMATION
We propose to train our models via joint variational inference for the last layer and MAP estimation of network weights (and noise covariance), yielding optimization problem
0''"", D" = arg max {2(0, n, Σ) + T (log ((0) + log p() - KL(g(E ( ())((())))
(24)
0.n.E
will write the three terms on the RHS (scaled by 1/T) as R(O, n. Σ). Reasonable priors for neural network weights have been discussed in several papers (Blundell et al., 2015; Pearce et al., 2020; Fortuin, 2022; Farqubar et al., 2020; Watson et al., 2020; Dusenberry et al., 2020; Nalisnick, 2018). In this work, use simple isotropic Gaussian priors which yields weight decay regularizer. While variational inference for the noise covariance possible, we choose (MAP) point estimation simplify the model. We use standard inverse- Wishart prior; ignoring terms that vanish gradient computation, we have likelihood
V N+N+
1 tr( tr(ME-)
log
logdet
(25)
2
where is N N, N are the degrees of freedom and M is the scale matrix. The terms v, M are hyperparameters that are fixed.
B.2 POST-TRAININ WITH VBLL LAYERS
In addition to jointly training the features and the last layer, potentially desirable several situations. For example, model may already be trained, and is useful to augment the model with uncertainty post-hoc. We propose to first train a model using standard network architecture and a standard loss function, yielding 0 and w² (or W in the multivariate case). Given these quantities, the last layer trained via
train them independently. This
" Σ arg max {((0" n7 E) T- P(E) KL ())l((E))) 7.E
(26)
Practically, one can initialize the mean of the variational last layer (in the regression of discriminative classification case) with the last layer point estimate wn from the first phase of training.
B.3 COLLAPSED VARIATIONAL INFERENCE FOR BAYESIAN NEURAL NETWORKS
Stochastic variational approximations to the posterior over the network weights have previously been used for Bayesian learning (Blundell et al., 2015). In this section, we discussion computation of variational posterior (0), following the SVI methodology as discussed previously. Whereas
16

---

Published as a conference paper at ICLR 2024
Algorithm 1 Variational BLL Training: Regression
Require: Training data D = [X,Y), variational posterior initialization v = (w, S), number of
train epochs N. minibatch size B, optimization algorithm opt(-).
1: for i = 1 to N do
2: Split dataset D in minibatches Dj = (X,, Yj), j 1...., 11/..
3:
for j 1 to [T/B] do
4:
0. w) + (()())) R(0, Σ) + T (KL(g(E n)]pp({ w)) - log logp(E))
[(0, n, Σ) < B ((x.))(X,,))) (- log (-
5:
6:
<- opt( VoC(O, n. VoR (,,)))
7:
7 V n, 1) V, n7 E))
8:
+ Σ (p 2) + VyR(O, 2)
9: end for
10: end for
approaches developed in the previous section were deterministic, SVI for all network weights is not possible via deterministic marginalization. Thus, computing
Egr0) [log p(y 2; 0)]
(27)
is typically approximated using Monte Carlo methods. In Blundell et al. (2015), the authors turn to the reparameterization gradient estimator (Kingma & Welling, 2014; Mohamed al., 2020) to enable the computation of the (Monte Carlo estimator the) gradient with respect to the parameters of the variational posterior. We could take a similar strategy for both & and 6, turning to sampling- based approximation. However, this sampling scheme yields both noisy gradient estimates and is expensive, as each sample corresponds to a full network evaluation. Our approach will instead marginalize the last layer and sample (some of) the other layers. This corresponds to Rao-Blackwellizann 1992; Blackwell, 1947) of the variational lower bound estimator, yielding lower variance gradient estimates.
(Rao,
We will choose posterior that factorizes over the (last layer) parameters and weights, q(E. 0 II)) = q(E ne (()) no). We also, in the discussion below, suppress dependence on >; in practice, we will turn to point estimation for this term. Note that further mean field factorizations for q(O n) are typically employed. For example, Blundell et al. (2015) factorize the posterior over all weights in the neural network. Given this posterior approximation, we have
log P( X)= 2 (m
[log p(Y X, g, 0)
(L(g(E) ne)ll((E)) KL(g(0) ne))lp(0) (28)
under the assumption that the prior P(E, 0) = (())p(0) and thus
1
T loy x)= 2 Ey(0)n)) [C(
(C ne)] 7 (L(EEE ne)llp(E)) - - <L((( ())))(()) (29)
for the lower bounds L developed in Section 3. Thus, algorithmically, we first compute the inner
expectation and then approximate the outer expectation with sampling-basee based estimator.
B.4 TRAINING
We now present our full training approach for the regression and classification settings. A detailed procedure for training the regression model with point features is shown in Algorithm 1. Generally, we will minimize the lower bounds we developed for each model. We note that L is sum over data; following Blundell et al. (2015), we compute an (unbiased) estimator č for this term with mini-batches.
The factorization of the ELBO over the data implics a mini-batch estimator for the gradient. Note that
T
1
Eq(elm) llog wr, &, 0) EgE (efttiimlloo [lo p(v, T
wx, E.]
(30)
where the outer expectation on the RHS is with respect to uniform distribution over ( 1,..., T. Note that this also holds for lower bound on the data likelihood, in the case of classification. We can construct a randomized estimator for this expectation based on sub-sampling the data, our case in mini-batches. For a mini-batch of B datapoints, this yields an estimator for the ELBO of the form
B
( n. Σ) B
Σ
Eyge m)) (glllgg(pl a. E, 0, E)]-
(31)
1=1
In the classification case, this may be an inequality. Note that in the limit of infinite training data (T -> 20) the weight on the KL term goes zero.
17

---

Published as a conference paper at ICLR 2024
We have trained VBLL models with both momentum SGD and AdamW (Loshchilov & Hutter, 2017). While both work effectively, they result in different uncertainty representations far from the data. The interaction of VBLLs with the stochastic regularization associated with different optimizers an important direction of future work. Practically, gradient clipping was necessary to stabilize late training, especially in the regression case. As the noise variance concentrates, gradient magnitude is highly sensitive to small perturbations to features, which can be rapidly destabilizing; gradient clipping was necessary and sufficient to prevent this destabilization. Beyond these details, training VBLL models did not differ from training normal models.
B.5 PREDICTION AND MONITORING
For prediction with VBLL models, we predict directly using the variational posterior, exploiting the conjugate prediction results described in Section 2. For all three VBLL models, training objective computation and prediction can be reduced from cubic to quadratic complexity (in the last layer input width) by careful parameterization and computation. The assumptions required to achieve quadratic complexity for the first two models are minor. However, for the generative classification model, diagonal covariances must be assumed. We diseuss complexity in the next section.
Training yields learned network weights O' (or variational posterior over these weights), noise covariance Σ", and last layer variational posterior parameters 17" To make predictions, there are two options. In the case of the regression and generative classification model, we may discard the variational posterior and leverage exact conjugacy. Under (Gaussian) distributional assumptions, exact posteriors may be computed with fixed features. However, exact last layer posteriors may be badly calibrated due to violation of distributional assumptions. Instead, we may make predictions under the variational posterior directly, under the assumption that q(E n²' = p(E X, Y), yielding
p x X, x E Ip(y x E. (.,0",""))
(32)
where (æt y) denote test point. For the discriminative classification model, only prediction under the variational posterior is possible, and in this model, sampling or an approximation (e.g. Laplace) may be used.
The generative classification case provides predicted class probability biases (the predicted probability of seeing a particular class before observing a label) through the Dirichlet posterior, In cases where a system designer believes there is likely to exist distributional shift between the training data and the evaluation conditions, predictions may be directly controlled by modifying this Dirichlet posterior.
For the variational feature approach, prediction can be done by sampling features and computing mixture distributions, yielding
ply w. X,Y R Egcor Eeigin F (p(y xr 3, 0,0,
(33)
K lpl x. E. 0x. Σ
E
(34)
K
for On sampled i.i.d. from the variational posterior. In the regression case, this averaging is straight- forward. For the classification cases, average pre-softmax or post-softmax. For example, in the case of generative classification, both
K
P(U 1 X,)
~
softmaxy | X,Y) + log Eg Eeerr y. 0x)l)
(35)
K
and
ply x, X) x softmax,,logp(( X, + log K Σ
E y, E, 0x)))
(36)
are valid Monte Carlo estimators for the predictive density, and the same holds for the discriminative classifier. In practice, we typically use the former (in which we directly average the sostmxx samples) due to the relative implementation simplicity, although the latter is necessary for some forms of out of distribution detection. Note that in the latter estimator,
log K
Σ D = LSE, (log xx) - log K k
(37)
for generic a. and log K vanishes in the softmax and my therefore be ignored, and where the use of log-sum-exp improves numerical stability.
18

---

Published as a conference paper at ICLR 2024
B.6 OUT OF DISTRIBUTION DETECTION
desirable feature of robust deep learning models is the ability to distinguish between in distribution and out of distribution (OOD) data. We use several metries for OOD detection with VBLL models. For the discriminative VBLL, we follow Liu et al. (2022) and use the maximum softmax proba- bility (Hendrycks & Gimpel, 2016) for an OOD measure. This is computed by sampling from the distribution over logits and passing these samples through the softmax, where they are averaged.
For the generative classification model,
use the feature density
p(a: X,Y) EE,.el((2, y E. E)
Σ
(38)
as an OOD measure. In the above, the expectation are with respect to the variational posteriors; for the MAP estimation case, this corresponds to direct evaluation.
In practice, we found post-training noise covariance calibration improved OOD detection performance for the G-VBLL model. More precisely, we aim replace shared diagonal Σ across all classes with Ey for each class. Our intuition is that while the Σ that is used in training prescriptive in the sense that it provides a model within which learning occurs the estimated per-class 2n are descriptive of the accuracy of modelling during training. Indeed, the training objective for the G- VBLL model is label (marginal) predictive likelihood, and so the training signal to model class feature densities highly accurately is weak.
Our calibration procedure as follows. First, we assume (MAP) point estimate for feature means My- For sufficiently large datasets Sy rapidly concentrates, so the impact of this assumption is relatively minor. For each class, we then compute the MAP noise covaraince Σn under the inverse- Wishart prior. Concretely, the mean under Gaussian prior N (it 2)
and known noise covariance Σ
is
By = (2 + Ty )(- Σ Φe + Eyy (ny)
(39)
T, - bE (.y))
(40)
where Ty is the number of class occurrences for class and where the sum is over all inputs in class y. For sufficiently large T and zero mean prior, this mean is approximately equal to the empirical
average Σ xz. Thus, taking jt = T Σ Ft, the noise covariance can be estimated as
1
S.y
(M+
( (,y)
(41)
Ty + +- N +1
which corresponds to the MAP posterior with a known mean, and where the sum is again over all inputs in class y.
We note that while our strategy of sequentially estimating two MAP estimates is relatively unso- phisticated, it is straightforward and yields good results, and is consistent for large datasets (under straightforward distributional assumptions). In the above, N corresponds to the dimension of the covariance matrix (as in (25)) and and M corresponds to the prior degrees of freedom and scale matrix, respectively. We found that this MAP covariances estimation outperformed the max likelihood covariance estimation as performed in Liu et al. (2022). Moreover, we note that both the empirical mean of the features for each class and the covariance can be recursively estimated in one pass over the data, and so the complexity of this step is OTT). Inspired by Ren et al. (2019; 2021), we subtract the log density under the feature prior as a normalization strategy, which also slightly improves performance.
While this post-training last layer posterior improves OOD performance, it is substantially over- concentrated for label prediction, yielding to dramatically over-confident predictions. It is an open question how to best estimate the last layer posterior to achieve both effective and calibrated label and OOD prediction.
C PARAMETERIZATION, COMPLEXITY, REGULARIZATION, AND HYPERPARAMETERS
In this section, we discuss how to parameterize each of the terms appearing in each type of VBLL. In each model, we use a "mixed" parameterization-iin contrast to the standard parameterization or natural parameterization. More precisely, we will parameterize the inverse noise covariance E-1 and the covariance of the variational posterior via Cholesky factorizations, and directly parameterize means W. μ.. In our (limited) comparisons of the performance of different parameterizations,
19

---

Published as a conference paper at ICLR 2024
found that our mixed parameterization performed equivalently (if slightly better) to the standard parameterization, and both performed better than natural parameterization. Interestingly, this stands in contrast to standard practice variational Gaussian process learning (Hensman et al., 2013), in which authors frequently aim to derive natural gradient optimization algorithms.
We will show that for each VBLL model, under a set of reasonable assumptions, complexity is at worst quadratic in the last layer width and linear in the output dimension. These complexity results enable use of VBLL models on problems with high input dimensionality and high output dimensionality. Moreover, our mini-batch gradient estimation training objective results in (standard) linear complexity of gradient estimation batch size, enabling training on much larger datasets than is possible with standard marginal likelihood objectives.
C.I REGRESSION COMPLEXITY
Our analysis will focus on the multivariate case, for which the univariate outputs are a special case. We directly parameterize the mean W E Nr X The covariances are parameterized via Cholesky decomposition to guarantee positive cemi-definttnnes; in particular we parameterize
Σ- = LL
L = Ld + diag(cxp(L))
(42)
5 PP
P = Pa + diag(exp(p)).
(43)
Where P. I. are lower triangular with positive diagonals, and thus L.a, Pa are lower triangular with zero diagonal, and vector L. p control diagonal elements.
Given these parameterizations, we show the complexity of each operation required for training is at most quadratic in Ng. The training objective has two terms: the log Gaussian density and the trace term. For the log density, we have
Tx-te=e_
(44)
for y WΦ. The term Le can be computed in O(N3) time. The second term is Φb S.,,r(E-1), for which of SΦ. can be computed in O(N2) time, and the trace term
tr(E-)) tr(LLT) ||]1
(45)
which can be computed in ((N)) time via squaring and summing the elements of L.
The remaining terms are the KL penalty on the variational posterior, and the inverse- inverse-Wishart prior on the noise covariance. Fixing a prior MN (w. I, S), the KL. penalty for the multivariate regression case is (ignoring constants)
det S det S
KL(g(E ()))((()) = 7((mn (rr((W W) T( W)S-1) + N2,r(s + N, log
(46)
We will fix an isotropic prior, S sT for 8 0. Thus, the first term
tr((W- W) (W - W) 1 W WIP
(47)
with complexity O(N, No), and the second term is
N,tr(S-1S)= - ur()
Nu
(48)
where the trace can again be computed as the squared Frobenius norm of the Cholesky factor of P. for complexity (N.)). The last term is
det
Ny log det.S
N₂ log s NylogdetS
(49)
where logdet. 2logdet(P) which is equal to the of the log diagonal elements, which can be computed in ((N)).
Finally, we have the inverse-Wishart noise covariance prior, which has terms logdetE and tr(ME-1) for scale matrix M. The log determinant term may be computed previously, with complexity ((₂)). Choosing scale matrix M mī, we have tr(ME--) mtr(E-1) which again is (NN3). Summing all of this up, we have the total complexity of VBLL computations as O + N2). +N2). which is equivalent to the complexity of standard matrix multiplication; thus, there is effectively zero added computational expense from the VBLL model compared to a standard network. The reader may easily verify that complexity of prediction is no greater than the training complexity in the regression model.
20

---

Published as a conference paper at ICLR 2024
Table 7: Time per batch on CIFAR-10 10 training.
Model
Run time (s)
% above DNN
DNN 0.321
0%
D-VBLL 0.338
5.2%
G-VBLL 0.364
13.4%
C.2 CLASSIFICATION COMPLEXITY
The complexity for the discriminative classification model follows from the regression model. We use the same parameterization, although we turn a diagonal noise covariance 2. The computation of the KL penalty is identical to the regression case. The only difference is that $ Sy must be computed for all classes y. yielding complexity ((N This term dominates the complexity of this model; however, further factorization of the covariance is straightforward and can reduce the practical complexity. To predict in these models, sampling realizations of the last layer must be done to sample logits. This sampling is straightforward to do using the Cholesky factorization of the covariance, and has quadratic complexity.
For the generative classification model, we are limited by the E + Sy term the log-sum-exp. As far as we are aware, there is no (practical) way to compute this term with quadratic complexity, or otherwise inexpensively compute this log density. Thus, in this paper we restrict E and S to diagonal matrices, which results in linear complexity in No for all operations in loss computation. Thus, under this approximate posterior, the complexity of the full training loss computation is O ((N which is equivalent to standard neural network models. This covariance structure is relatively restrictive, and improvements may results from sparse covariance structures.
Concretely, we compare the run time of one step of training across baseline DNN, and both flavors of VBLL. We compare these models on CIFAR- 10 training on a NVIDIA T4 GPU, with the wide ResNet encoder used in the rest of the classification experiments. The results are shown Table 7. We note that our VBLL implementations are not carefully optimized, and so these slowdowns are an upper bound on the possible slowdown.
C.3 COMPLEXITY OF COMPARABLE BASELINES
There are a set of baseline methods that are similar to VBLLS but often have different complexity. As discussed throughout the paper. training BLL models by exploiting exact conjugacy (or exactly computing the marginal likelihood) requires iterating over the full training set. yielding lincar complexity in the size of the dataset. This almost always makes standard marginal likelihood training intractable. More directly comparable is SNGP (Liu et al., 2022), which also exploits exact conjugacy (or approximation thereof for classification) but only computes the last layer covariance once per epoch. This amortizes the cost of iteration over the full dataset. In practice, they use an exponential moving average estimate of the covariance. which removes the need to load the data multiple times per epoch. However, this covariance must still be computed and inverted, which has cubic complexity in the last layer dimension. Last layer Laplace (Daxberger et al., 2021a) methods, similarly, require a pass over the full dataset and must invert a dense covariance matrix, yielding cubic complexity. However, this is only done as a post-processing step for a trained model.
C.4 HYPERPARAMETERS
VBLL models introduce a small number of hyperparameters over standard network training. First, standard hyperparameters may need to be modified for VBLI. models. For example, we found longer training runs resulted in slightly improved calibration, but we believe further investigation of learning rate schedules is necessary. For MAP features estimation, we use standard weight decay regularization values.
The main novel hyperparameters introduced by the VBLL model are those associated with priors. In particular, the last layer mean prior (defined by a mean and variance; in the regression case, these are written w. 5) must be chosen. Practically, it is common to normalize outputs to have isotropic Gaussian distributions for regression, and thus we have found w = 0 and $ = I yield reasonable if diffuse prior. For the classification case, we found these values similarly induce reasonable epistemic uncertainty over the predictive categorical distribution.
The other novel hyperparameters are those associated with the noise covariance inverse- Wiishart prior, the degrees of freedom v and the scale matrix M. For all experiments, we fix the scale matrix as a scalar multiple of the identity matrix, M = ml. In our regression experiments we fix these to be (1, 1), and find good resulting performance, but further investigation is possible. In the classification case and in particular the generative classification case-these parameters control the degree of
21

---

Published as a conference paper at ICLR 2024
Regularization
Regularization
6-
4-
2
2-
0-
2 0.0
0.5
1.5
2.0
2.5
3.0
0.5
1.0
1.5
2.0
2.5
exp(2pb)
expl2p))
Figure 3: Weight decay (left) and KL/Inverse- Wishhrt regularizers (right) plotted versus exp (ps) (which corresponds the diagonal element of the covariance matrix). Different curves show varying weight decay strength and varying a term in (56), with b
concentration in the feature space, and thus must be more carefully selected (and often co-selected with the weight decay strength).
C.5 UNDERSTANDING PRIOR REGULARIZERS
In this subsection we investigate the regularization effects of the prior (and KL) terms, and com- pare them to standard regularizers such as weight decay. Note that naive weight decay on these parameterizations would correspond to additional loss terms of the form
(1I + |1012 + ||aal| + lPl22)
(50)
The loss terms resulting from our chosen priors in the regression case (and dropping terms with gradient) are
- logg = (mir(>=1) - illogdet>
(51)
2 W N 1r N, NoJlgdet5))
KI(g(E m))lr(())
(52)
for = 4 +++N++ 11 note that (other than the weight decay-like term on W) both covariance regularizers are of the form
atr(M (M) blogdet(As).
(53)
for constants a, and matrix M. Given our Cholesky parameterization,
tr(--1)= = |PP + Σ exp(2pk)
(54)
logdetZ-1
2pk
(55)
and similarly for S. Thus, the regularization of the off-diagonal covariance terms again corresponds simply to weight decay, whereas the diagonal elements of both covariance matrices have regularizers of the form
Σ ((eexp(2p)) - 2bpx).
(56)
Note that this function is convex. This function (inside the summation) is visualized for varying (compared to weight decay) in Figure 3. Our regularization terms provide substantially more control over the minimizing value, and thus more control over predictive variance. However, compared to weight decay, our regularizers vary in scale substantially more which may lead to difficulties trading off regularization terms with other loss terms.
To counteract this relative lack of interpretability of our hyperparameters, we propose an alternate representation of these values. We rewrite the regularization function as
(exp(2p=) - p
(57)
where a corresponds to a scale term, and b/a controls the location of the minimum. We may specify a desired predictive variance, which can be mapped to the minimum of the regularization function. Concretely, given some target variance element 8 exp(2pt) (for all k), we choose
ša
(58)
22

---

Published as a conference paper at ICLR 2024
Regularization
Regularzzaton
6-
4-
2-
2
0-
0
0.0
0.5
1.5
2.0
2.5
3.0
0.5
1.0
1.5
2.0
2.5
exp(2ps)
expl2pt)
desired predictive variance 8, with a 100. Right: sweeping over regularization scale with fixed desired predictive variance 8 = 1. Note that all functions asymptote exp(2ps) = In these figures, the curves have been vertically shifted to achieve a minimum at zero; this vertical shift does not impact regularization.
Figure 4: Sweeping
modified hyperparameter representation. Left: sweeping
which assures that the minimum of (56) achieved when p; 1 log 8 for all Sweeps over the hyperparameters (a, 8) presented in Figure 4.
Given this transformation between hyperparameters, these alternate hyperparameters in VBLL models. The original hyperparameters for the model, as described earlier in this section, are the prior last layer covariance scale s, the scale matrix for the noise covariance prior m, the degrees of freedom D. Additionally, it is common is Bayesian deep learning to scale down the KL penalty, and we write this factor as A. Our alternate hyperparameters
can now be concrete in how specify
target (diagonal) values, i >> 0 and p and scale parameters ag and ag 0. The mapping between these hyperparameters is:
+ D
m f- a
(59)
pTas N
las.
(60)
X as (perhaps naively) theoretically justified in variational inference, then og is correspond- ingly fixed.
D EXPERIMENTAL DETAILS
This section contains details about the experiments the body the paper. We note that for highlighting in the tables in the paper body, if a single-pass method (in the upper half of each table) is the best performing in metric, that result is highlighted. If the best performing is multi-pass, pasti, we highlight both the best multi-pass and single-past method in the column. We believe that this important, as many applications required single-pass methods and thus multi-pass results irrelevant.
D.1 METRICS
For regression experiments, we report the predictive negative log likelihood (NLL) of test data, which can be computed in closed form for point feature estimates. We also report the root mean squared error (RMSE), standard metric for regression. For classification, in addition to the negative log likelihood, we also report predictive accuracy (based on standard argmax of the predictive distribution), and expected calibration error (ECE), which measures how the model's subjective predictive uncertainty agrees with predictive error. Finally, we also investigate out of distribution detection performance, a standard evaluation scheme for robust and probabilistic machine learning (Liu et al., 2022). We compute the area under the ROC curve (AUC) for near-OOD and far-OOD datasets, which is discussed in more detail later in this section.
D.2 BASELINES
We distinguish baselines between single-pass and multi culti-paass models, which we show in upper and lower segments of each table, respectively. Single-pass methods require only a single network evalua- tion, and we compare VBLLs with MAP feature estimation to these models. Multi-pass methods require several network evaluations, and includes variational methods like Bayes-by-backye (which refer to BBB) (Blundell et al., 2015), ensembles (Lakshminarayanall al., 2017), Bayesian
23

---

Published as a conference paper at ICLR 2024
dropout (Gal& & Ghahramani, 2016) and stochastic weight averaging Gausssinn (SWAG) (Maddox et al., 2019).
Within regression, we compare to models which exploit exact conjugacy, including Bayesian last layer models (GBLL and LDGBLL (Watson et al., 2021)) and RBF kernel Gaussian processes. We note that these methods require computing full marginal likelihood and are thus difficult to scale to large training sets. We also compare to MAP learning, in which a full network is trained via MAP estimation, and Bayesian last layer is fit to these fixed features (Snoek et al., 2015). Within classification, we primarily compare to standard networks (DNN), as these output a distribution over labels and thus can be directly compared to our model. We also compare to SNGP (Liu et al., 2022) and last layer Laplace-basee SNGP aims to approximate deep kernel GPs (Wilson et al., 2016b), and Laplace methods compute a last layer approximate posterior after training. We note that in contrast to SNGP (Liu et al., 2022), we do not modify standard neural network backbone, such as including spectral normalization, adding residual connections, or using sinusoidal nonlinearities. Both SNGP and last layer Laplace methods require a pass over the full dataset to fit the last layer distribution; in contrast, our method maintains a last layer distribution during training, which may be useful for e.g. active learning. We do not evaluate Laplace methods in regression as they are nearly identical to the MAP model.
methods (Daxherger et al., 2021a), which are similar last layer models.
D.3 TOY EXPERIMENTS
Figure 1 contains simple visualizations for the regression model and the generative VBLL model. In particular, the regression model shows predictions with variational feature learning (with KL weight of 1.0) on cubic function with gap in the data. This dataset consisted of 100 points sampled in [-4.-2] U [2. 4]. with a noise standard deviation of 0.1. The model consisted of a two hidden-layer MLP of width 128, trained for 1000 epochs with a batch size of 32, with stochastic gradient descent with momentum, with a learning rate of 3 10 zero weight decay, and momentum beta parameters of 0.9. These values were arbitrarily chosen, although the choice of SGDM versus Adam (Kingma & Ba, 2015) does make a difference on prediction far from the data, Gradient clipping with a maximum magnitude of 2.0 was used. The DOF and scale parameters were both set to 1.0
For the classification problem. we used the scikit-learn (Pedregosa et al., 2011) implementation of the half moon dataset, with 1000 data points and a noise standard deviation of 0.2. We trained G-VBLL model with residuul-structureed MLP of width 128 (each hidden layer added to the layer input). This model was trained with SGDM with learning rate 3- 10-2, momentum beta 0.9, and weight decay 10-4, 4, for 100 epochs and with a batch size of 32. The DOF parameter was 128, and the scale parameter was 1.0.
D.4 REGRESSION
Our UCI experiments closely follow Watson et al. (2021), and we compare directly to their baselines. For VBLLs, we used (0, 1) last layer mean prior and W-1(1, 1) noise covariance prior. For all experiments, we use the same MLP used in Watson et al. (2021) consisting of two layers of 50 hidden units each (not counting the last layer). For all datasets we matched Watson et al. (2021) and used a batch size of 32, other than the POWER dataset for which we used a batch size of 256 to accelerate training. For all datasets we normalize inputs (using the training set statistics) and subtract the training set means for the outputs. We did not re-scale the output magnitudes, to retain comparability of NLLs. We note that the extent to which outputs were normalized in Watson et al. (2021) is unclear. However, they make the parameters of their prior Icarnable, which can have a similar effect to centering the outputs, and so we believe our output centering is reasonable. All results shown in the body of the paper are for leaky ReL.U activations. For all experiments, fixed learning rate of 0.001 was used with the AdamW optimizer (Loshchilov & Hutter, 2017). A default weight decay of 0.01 was used for all experiments. We clipped gradients with a max magnitude of 1.0.
For all deterministic feature experiments, we ran 20 seeds. For each seed, we split the data in to train/val/test sets (0.72/0. 12/00.18/// I of the data respectively). We train on the train set and monitor performance on the validation set to choose total number of epochs. In contrast to Watson et al. (2021) who compute validation performance for every epoch, (predictive NLL) every 10 epochs (note that the datasets are small and typically train for hundred of epochs). After choosing a number of epochs, we train on the combined training and validation set and evaluate performance on the test set. We use a max number of epochs shown in Table 8, which were large enough to not be reached but often lower than those used in Watson et al. (2021).
compute validation performance
For our BBB feature models, we ran 10 seeds with a similar procedure to the above. We follow Watson et al. (2021) and use (0, 4/ /Tim) for each weight (where nin denotes the layer input
24

---

Published as a conference paper at ICLR 2024
Features
BOSTON
CONCRETE
ENERGY
POWER
WINE
YACHT
MAP
3000
3000
2000
3000
1000
2000
Variational
10000
10000
10000
10000
10000
10000
Table 8: Maximum number of epochs for each set of features and each UCT dataset.
Feature KL
Weight
LL KI. Weight
MAP
50
5
0.5
1.0
0.160
0.266
0.281
0.282
0.1
0.162
0.266
0.286
0.272
0.01
0.168
0.268
0.268
0.280
0.001
0.160
0.267
0.272
0.276
Table 9: CIFAR-I0 NLL for varying values of KL weights, for both the last layer and the feature weighting in variational feature learning.
width). Validation performance was monitored every 100 epochs, and 10 weight samples were used to compute the validation predictive likelihood and choose full training number of epochs.
D.5 IMAGE CLASSIFICATION
All classification experiments utilize the Wide ResNet-28-10 (WRN-28-10) backbone network architecture. Hyperparameters are similar to those proposed by Zagoruyko Komodakis (2016). Unlike the original implementation WRN, we do not employ Nesterov momentum and we fully decay an initial learning rate of 0.1 according to a Cosine Annealing schedule instead of a stepped decay schedule. Gradients are clipped with a maximum magnitude of 2.0 and we impose last layer KL weight of 1.0. We All classification results are reported across 3 seeds and use the standard WRN data-augmentations proposed by (Zagoruyko & Komodakis, 2016). For the deterministic feature experiments, we train each model for 300 epochs.
The BBB backbone-base models utilize the same WRN architecture and are primarily deterministic. The BBB models implement single final Bayesian linear layer with a prior distribution of (0, 0.01). Each BBB-based model used 10 weight samples for test set evaluation. This operation is relatively cheap when compared to fully stochastic network because the intermediate features are cached prior to the final Bayesian linear layer weight sampling and computation. All BBB are trained for 400 epochs and we impose a last layer KL weight of 1.0 and a feature KL weight of 0.5 the VBLL-BBB and DBLL-BBB models. The BBB baseline model utilized a feature KL weight of 50.
D.6 SENTIMENT CLASSIFICATION WITH LLM FEATURES
We perform sentiment classification experiments utilizing features extracted from a pre-trained OPT-175B (Zhang et al., 2022) model on the IMDB Sentiment Classification dataset (Maas et al., 2011). 1). We compare our G-VBLL and D- D-VBLL models with an MLP baseline. The IMDB dataset text-based binary classification task in which inputs are polarized movie reviews and outputs are positive and negative labels. Text embeddings are extracted from the OPT-1755 model for each sample as the output of the last model layer for the final token in the text sequence. This results R 12288 for each sample. In all cases,
in sequence embedding, prior to the classification head. To understand the impact of training dataset size on performance, all experiments are performed at multiple training dataset scales. The IMDB dataset is sampled iid. to construct training datasets with 10, 100, 1000 samples alongside the standard 25,000 sample training split. We train models at all dataset scales and report across 3 seeds. The Adam optimizer is used for all models. Hyperparameters such learning rate, weight decay were tuned across both the 10 sample and full dataset scales.
utilize two linear layers
D.7 WHEEL BANDIT
We match the experimental settings of Riquelme et al. (2018). In particular, we use batch size of 512, learning rate of 3e 3, and train for 80000 steps total. We perform 20 steps in the environment per phase of updating, and perform 100 gradient steps when updating. We use gradient clipping norm of 1.0. We use the same network architecture as baselines, an MLP with widths (100, 100, 5). where the last layer is VBLI. The VBL.I. hyperparameters were set to defaults: the degrees of freedom and the scale in the Wishart prior are set to 1, and the prior scale was also set to 1.
25

---

Published as a conference paper at ICLR 2024
0.01
0.01
10a
107
1.0
1.0
10a
10.0
10a
10.0
Covariance Norm
Covanance
10a
10
10a
10
10
10-2
10
250
500
750
1000
2000
200
600
1000
Epoch
Epoch
Figure 5: Sweeping over the E location parameter for UCI datasets Energy (left) and Wine (right). The dotted colored lines correspond to Σ values over the course training, and solid colored lines correspond the Frobenius norm of S. The black dotted lines correspond to target > values. The scale hyparparameter was large these experiments to illustrate the ability to effectively control noise covariance. Note that for very small E- the impcat of the predictive loss limits the degree which realized noise covariance matches the goal value; this trade-off is controlled by scale parameters.
Covariance Norm
10a
Covariance
10c
0.1
1.0
10-1
10
0.1
1.0
10
10"7
250
500
750
1000
1250 1500 1750 2000
200
400
600
800
1000
Epoch
Epoch
Figure 6: Sweeping over the S location parameter for UCI datasets Energy (left) and Wine (right). Again, dotted colored lines correspond to Σ values over the course training, solid colored lines correspond to the Frobenius norm of S, and black dotted lines correspond to target diagonal S values. Note that the Frobenius norm of S in all cases is higher than the target due to the off-diagonal elements, but the realized covariance can be well controlled.
E HYPERPARAMETER STUDIES AND ABLATIONS
KL weight. We additionally explore the NLL sensitivity of the DBLL and DBLL-BBB models to various KL weighting configurations. In Table 9, we sweep across orders of magnitude for both the last layer and feature KL weighting parameters.
Location and scale hyperparameters. We investigate our hyperparameter reformulation on UCI datasets in Figures 5 - 7. In particular, we vary each of the location and scale parameters and show that we can effectively control the quantities of interest. In particular, Figures and show varying the location hyperparameter for each covariance matrix S with high scale hyperparameter, enabling fine-grained control over realized values. In practice, this degree of direct control over realized model values is not desirable, and these plots only illustrate that such degree of control possible. In Figure 7, we vary the scale parameter for Σ and show that it effectively controls the strength with which Σ is regularized. With naive hyperparameter selection, interaction between scale and location parameters would require careful planning to control regularization scale independently of location, whereas reformulation enables direct control of scale.
26

---

Published as a conference paper at ICLR 2024
10a
10
100
100
1000
1000
104
Covariance Norm
10a
Covarance
10
10a
10a
10-3
10
250
500
750
1000
1500
2000
200
400
600
800
1000
Epoch
Epoch
Figure 7:
Sweeping over the Σ scale parameter for UCI datasets Energy (left) and Wine (right).
Again, dotted colored lines correspond to E-1 values, solid colored lines correspond to the Frobenius norm of S, and the black dotted line corresponds to location hyperparameter >' value of 1. Note that by varying the scale hyperparameter, the strength of the regularization is varied without changing the target value, which a result of our hyperparameter reformulation.
F
PROOFS AND FURTHER THEORETICAL RESULTS
F.1 HELPER RESULTS
Our first result builds on results from the variational Gaussian process literature (Titsias, 2009; Hensman al., 2013).
Lemma 4. Let q(v) S) and ply X N and S,E R RMM Then
x1) N(X Σ) with y R.V p. B,M,
Eg(u) Dog p(y X, X,)+] = logrly X, -trE--XXX"
(61)
Proof. We have
1
Egin) [log
X =
E d(m) (lgdellgd)((2x1) (y - X11) E 2 XP)
(62)
2
1
2 (logdet(2x + Eq(n)((y X14) E - Xp)))
(63)
2 (logdet(2 + - (y-Xp) X (y X + tr -
=
(64)
where the last line follows from the fact that X/u ~ N XI, XSXT), The first two terms form the desired log density.
Based on this result, we can state straightforward corollary for generative classification.
Corollary 1. Let q(μ)) = Na, S) and p(y μ)) = NN,, 2) with y, p, p E R.N S, NNN
Then
E Doggly [+)] = log ( (1(E-1)).
(65)
Proof.
This result follows from Lemma 4 by simply choosing X = I.
We can also present variant for multivariate classification.
Corollary 2. Let q(W) MN(WW, S) and p(y xx, W) N(W with RM W.W MXN, RN S R N*N; and E E R M Then
(ogg(y) x. W)] = logp(u x. W)
2 S
(66)
27

---

Published as a conference paper at ICLR 2024
Proof. Our proof closely follows that of Lemma 4. Expanding the likelihood in the expectation, have
(g (eww)logp(s) ply æ, W)] p(y æ,, W) Ewlæ (xr (W - -- W)
(67)
Leveraging the matrix normal identity
Ew-MN(W,v.n
AW] = [WTnp]=Wwwwwwwo.
(tt((AVV) W AW
(68)
and the fact that W- W (O 1, S), have
EE(W W) E S-'(WWWW)))W - [W)] = Str(E=
(69)
which completes the proof.
Lemma 5, Let p(æ (4) = (H, E), and let μ [p, S). Then, (p [ (i, E 5).
(70)
Proof. We build upon Jacobson (1973) and note
del exp(-51 1
s))xe T Σ _-122)]]
(--
p))
V dot(s-1+ + +--1)
(71)
Note, by Woodbury's 's identity
S (S (I-1 --==
(S + 2;-
(72)
Let == = - ja, then 2 ~ N(x p. S). We then have
Elp(2 pd] =Elexp( 11 - 11/2-11 + logdet (0ggeet2+=E-11 (2+
(73)
2
Efexp(- 2²
"T (8-1z))] 2)] exp( 2
(74)
For the expectation we apply (71). We simplify the determinant term of (71) as
del(S-1)
V d =exp(- - logdet(T + +S=-!)
(75)
Combining, we have
1 (1 - 1 (s+ + logdet(7 + +SE-1)
Elexp(- 2 We have two log determinant terms, from (74) and the above. We can combine them as
- = exp(-
(76)
logdet(2*E
2 logdet(1 + +SS- = (logdet(
Σ) + logdet( + 1)
(77)
2
2
2n
logdet((
Σ) + + 1)
(78)
2n
1
1
logdet(
Σ+
5)
(79)
2T
2m
Combining all terms completes the proof.
F.2 PROOF OF THEOREM
Theorem 1. Let q(E () = N(w, S) denote the variational posterior for the BLL model defined in
Section 2.1. Then, (12) holds with
T
C(O, E)
(1ogN(en
w Φr. Σ) 20. S0.--1)
TT..
(80)
T
Proof. First,
logp(Y () log EpY X, E. 0)]
(81)
X, E. p(E)
Eglem)) (
(82)
q(E m)
> Eg m) blog p(Y
X, E 0)] - KL(g(E) (p)))
(83)
Egten dogp(v1 æ E. 0)] - KL(G(E ())p(()))
(84)
28

---

Published as a conference paper at ICLR 2024
Note that the first term in the last line is the log of Normal distribution. Applying Lemma have
(v(e))) logp(ue De; 0)] = log p(y: E. 2 20, SΦ.
(85)
which completes the proof.
We can also state the following corollary for the multivariate case.
Corollary 3. Let qiE |n) = -M,WW, I,S) denote the variational posterior for the multivariate BLL model defined in Appendix A. Then, (12) holds with
L(
1 Σ T
[og.N(y W o. Σ) 2 (()
(86)
Proof. The proof follows the proof of Theorem 1, applying Corollary instead Lemma
F.3 PROOF OF THEOREM 2
Theorem 2. Let q(W | 1) = IT
k=1 (ww,, St) denote the variational posterior for the discrimina-
tive classification model defined in Section 2.2. Then, (12) holds with
W Φ: + [ SKO. + oi)
1
[( n. Σ)
y, Wo LSE
(87)
T
Proof. We construct an ELBO via
logp(Y x, 0) log Ep(e) (p(Y X, [)
(88)
E pY PY X, 0,0 KL(g(E (n)lp(())))))
(89)
(em y log softmax,
0 E))] KL(((E ()))((())
(90)
Expanding the log-softmax term, we have
Ey(Em)) [y. log softmax, (log p(=,, y 0, [0ε))]
(91)
Egggin) [3/ log poz,( y (){ Egggin) [LSE,,[og p(æ;, y 0, 0.E)
As previously, under the variational posterior these likelihoods factorize across the data. The first term may be directly evaluated, yielding
Eg(gin) (logg((x,)) y 0, [))] = Egggin) [ (
w
(92)
The second term (containing the (og-sum-xpp) can not be computed exactly, and will bound this term for both the discriminative and generative classifiers. Via Jensen's 's inequality, have
- Egg [LSEy p(z;
0, () 2 - log V
Eytelm) [exp(log p(x,, e. y 0 (.E))]
(93)
the case of the discriminative model, we follow Blei Lafferty (2007) and note that for each row
Ew (exp(ut db ex)] exp(wi Φ: (Φ) Se + +o)))
(94)
which relies on assumed independence of rows of W (although relaxation of this assumption is possible). Combining these results yields lower bound on the ELBO, which is itself a lower bound on the marginal likelihood.
F.4 PROOF OF THEOREM
Theorem 3. Let g(p. In) = N( denote the variational posterior over class em-
beddings for the generative classification model defined in Section 2.3. Let p(p Y) = Dir(a) denote the exact Dirichlet posterior over class probabilities, with a denoting the Dirichlet posterior concentration parameters. Then, (12) holds with
C(O, m. Σ)
(log Py. 2) Str(I-'Sy)) (b (()()))) + log as (95)
2
LSEx(log.A ((ollloo a E+ Sk) + log ass))
where (.) is the digamma function and where a,
Σx Ock-
29

---

Published as a conference paper at ICLR 2024
Proof. Note that
lY 0) 2 EE m) (Y X x, 0)] KIL((E m))lp(E)
(96)
where
Ey(elm) X, 0)] Ey(Ein) [logp(X Y, 0, logp(X 0, <)] + Eg(eiv) (lg
[)]
(97)
All of these terms factorize over the data, previously. We first note that for the last term, Epllogp(3: p)] = ((a))) 15
(98)
where a correspond to posterior Dirichlet concentration parameters and (-)-) denotes the digamma function. The first term in (97) is the embedding likelihood; we can compute this expectation of the log likelihood via Corollary I.
The second term in (97) is less straight-forwand Note that
E| Eloogp(z; 0, <)] = Elloo Ep(az E)p,)) 0, 0,E)]
(99)
y/
which can be written as a log-sum-exp of log joint likelihood. exchange the log and sum, and note
We will again apply Jensen's 's to
Ellog p(x, 0, - E [log Σ p( y. (p(y) (,
(100)
log El > p(a; y, 0. () 0.
(101)
log Σ En (o(a; My. 0))E p)]
(102)
-LSE [log (p p)] + log Fx Ip(are My. 0)]]
(103)
where the second line follows from Jensen's, and the third line follows from the structure of the variational posterior. We may apply
log p(y 0, p)] = log Cty - llog Σag.
ag. y/
(104)
a standard result from Cartgoiet-Categorrial marginalization. The second term (104) (the sum over concentration parameters) is equivalent for all classes y. and thus can be pulled out of the log-sum-exp (due to the equivariance of this function under shifts) where it cancels the same third term in (97).
To compute the second expectation in (103), apply Lemma 5. Combining all terms completes the proof.
30


Reviews:
[{"review_id": "5MLJrut1uF", "reviewer": "Reviewer_tM8T", "summary": "The authors introduce variational Bayesian last layers as a novel approach for approximate inference in Bayesian deep learning models. The main contribution is three-fold: (i) following the current trend in Bayesian deep learning the authors propose to use a variational approximation to the last-layer posterior, (ii) the authors introduce closed-form bounds on the ELBO for different likelihood functions, (iii) the authors show that the simple approach can result in improved performance for regression and some classification tasks.\n\n--\n\nI have adjusted my score based on the author's response.", "questions": "1. Eq. 14 uses a rather loose bound, is it possible that this is the reason the approach underperforms in the classification settings compared to the regression setting? If so, is there any way to obtain a tighter bound?\n2. How does the method perform if it is used only as a post-hoc approach, meaning, without adaptation of the feature map? In large-scale applications, this is a particularly relevant setting and the proposed method could be a promising plug-in replacement.\n3. From what I understand T is never actually used. Is this correct?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and easy to follow in most parts. Moreover, the work is well-motivated and I enjoyed that the authors brought back old ideas to the BDL community, e.g., using the discriminant analysis as a likelihood model.\n2. I believe the exposition of the method is well done in most places, though slightly dense here and there, and helped in understanding the general idea of the proposed method. Moreover, I believe that the method is correct and an interesting contribution to the field. I think it is important to see more work on deterministic approaches to uncertainty quantification in deep learning. \n3. The experimental section shows promising results, especially in the case of regression.", "weaknesses": "Overall: My main concern with the paper is the weak empirical evaluation and limited novelty of the work, that is, it seems it is essentially an application of known techniques to the special case of last-layer posteriors.\n\nComments:\n1. Section 2.4 lists various related works, which I believe the author claims to optimize the log marginal via gradient descent. I have not checked every citation, but it appears to me that this statement is false for at least a subset of the cited papers. It might be good to revise the exposition.\n2. Eq 12 is some weighted ELBO, weighted with T for the purpose of generality, according to the authors. However, T never seems to be used later and makes the connection to the common ELBO less transparent. I believe the paper would improve in clarity if T is dropped.\n3. Section 3.4 is very dense and it could help the reader if this section is improved in its presentation. \n4. For the experiments, I would have expected assessments under distributional shift, a comparison to recent deterministic approaches (e.g., Zeng et al 2023 or Dhawan et al 2023), and a large-scale application of the approach as it acts on the last-layer only and should be applicable in more realistic scenarios (e.g., ImageNet).\n\n\nMinor:\n- Page 3, Eq 11 cites \"Harrison et al 2018\", which I looked up but didn't find any relevant content that would discuss the use of the marginal in Bayesian deep learning as a standard objective. What is the reason for the citation?\n\nCitations:\n- [Zeng et al 2023] Zeng et al, \"Collapsed Inference for Bayesian Deep Learning\", NeurIPS 2023.\n- [Dhawan et al 2023] Dhawan et al, \"Efficient Parametric Approximations of Neural Network Function Space Distance\", ICML 2023.", "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": "QLTuoubpKN", "reviewer": "Reviewer_vugz", "summary": "The paper proposes a novel way to learn Bayesian last layer (BLL) networks by deriving a sampling-free approach for optimizing normal, categorical, and generative classification likelihoods,\nby relying on variational inference techniques. \nThe approach is then evaluated on a series of regression and classification tasks against a range of baselines.\n\n\n_____\n_Edit: Given the improved presentation and evaluation, I increased my score._", "questions": "- Why was this specific subset of six UCI data sets chosen? The original work by Hernández-Lobato and Adams (2015), who introduced this set of experiments had ten, and even Watson et al. (2021) who the authors cite as relying on for their setup used ~~seven~~ different sets. _(PostRebuttal Edit: I misread the reference, Watson et al. use the full set of experiments.)_\n- Can the authors provide further results on the empirical runtime of the proposed approach, not just a theoretical one?\n- What was the principle according to which average numbers are bolded? E.g., in Energy RMSE a huge range of means is bold (from 0.39 to 0.47), but 0.43 is missing;  CIFAR-100 AUC has the same pattern, huge range, some missing, etc. \n- (very minor) What is the irony in BLL methods being popular (Sec 2.4)?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "BLL networks are an interesting approach to solve the scalability problem Bayesian neural networks tend to suffer from. \nThe paper introduces another variation to this family of approaches that is relatively straightforward, easy to understand, and implement.\nThe method is properly evaluated as the number of experimental setups is reasonably extensive both with respect to architectures and experimental tasks.", "weaknesses": "Straight-forward contributions can be seen both as a strength and as a weakness depending on the situation. \nThey are a strength if they are an easy solution to a complex problem that might not improve upon current approaches in all situations, but most. \nThey are a weakness if they do not provide a clear theoretical benefit above current approaches and also come without clear performance improvements.\nFor me, the results point to the latter case as they are rather mixed despite some strong wordings of the authors in their claims. \n\nThe abstract promises \"improve[d] predictive accuracy, calibration, and out of distribution detection over baselines.\", similarly in the contributions, and conclusion \nparts of the paper. Even more, the method not only improves, but it also performs \"extremely strong\" and \"exceptionally strong\".  \nThese are some exceptionally strong statements given the actual performance.  \n\nFocusing on each of the experiments in turn. The first problem is the presentation, e.g., what does a bold number mean (see question below)?\n\n_Regression Experiments._ Of \nthe six data sets (see question below on this number) the proposal improves on two, slightly on one, equally on two (although better than neural net-based baselines), and worse than most of its baselines in the final one. Calling this \"strong performance\" is rather misleading. Two additional, though potentially minor, problems are that all of the baselines are simply cited from prior work (Watson et al., 2021). Given the wide performance variations between different train/test splits that can be observed for various UCI data sets the results are not entirely trustworthy. (Note that the reported error intervals are most likely standard errors, as is common on UCI, instead of standard deviations. But which they are is never specified.)\nSecondly, the authors acknowledge in the appendix that there might be differences in the way the training data is normalized compared to the cited results. \n(Whether these problems strongly influence the results, or bias them in favor or against VBLL is unclear.)   \n\n_Classification._\nWhile \"Extremely strong accuracy results\" are mentioned, it just performs as well or worse than competitive baselines like Laplace or Ensembles. The same for ECE, NLL, OOD detection.  where \"exceptionally strong performance\" is claimed.\n\nThe method is somewhat simpler than baselines, but it lacks a convincing argument for why this should matter. As the authors advertise this simplicity, there should be additional results on practical runtime improvements compared to the baselines to provide some evidence for the claim that a reader should use this approach. \n\n### Minor\n- Dropout is first claimed to have a \"substantially higher computational cost\" (Sec 1) and appears later as a \"comparatively inexpensive\" method (Sec 4).  Additional forward passes at test time are indeed rather inexpensive instead of a high computational cost.\n- When submitting to ICLR please make sure that you follow the ICLR style guide. E.g., Table captions belong above tables, not below.\n\n### Typos\n- Sec 5.2 first par: \"We refer to this model **as** G-VBLL-MAP...\"\n- (6): $y|\\rho$", "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": "mLGkQTisNv", "reviewer": "Reviewer_tT7z", "summary": "This paper derives efficient optimization objectives of the posterior distribution over the parameters of the last layer of neural networks for common machine learning applications such as classification, regression, and generative classification. Theoretically, their variational inference objective functions are derived in closed form and therefore enjoy the property of not requiring sampling - meaning the cost of enabling uncertainty quantification for a broad class of deep learning architectures is marginal (parameters other than the last layer are learned by maximum a posteriori). Experimentally, they validate these novel variational inference algorithms using standard benchmarks from UCI, a large language model used for sentiment analysis, and an image classification problem.", "questions": "For the Resnet image recognition and sentiment analysis experiments, what was the additional compute required (or time taken per iteration, if available)? The sample-efficiency is great, and understanding the practical overhead rather than theoretical complexity would be great for larger models that are in broad use.", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This is an excellent paper and a significant contribution - well done! The authors make clear how they build on the existing literature in Bayesian deep learning to create a novel advance that is practical and easy to implement. This is significant and should enable more work to push the frontier of the \"best of both worlds\", with neural networks serving as function approximators and Bayesian methods enabling sample-efficiency and quantification of uncertainty that is required for practical deployment of deep learning.", "weaknesses": "Visualizations of how tight or loose the bounds in the main text could help build more intuition; comparisons in terms of speed or efficiency to variational inference algorithms that do require sampling (such as Monte-Carlo objectives like VIMCO) could also help guide practitioners in making the correct trade-off depending on FLOPs of compute available versus the required accuracy of posterior approximation/uncertainty quantification.", "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": "2zHIwsqRN5", "reviewer": "Reviewer_b26R", "summary": "This paper introduces Variational Bayesian Last layers, a technique to perform uncertainty estimation in standard neural network architectures. \nThe method performs bayesian learning only for the last layer in the neural networks using Variational Inference.\nThis results in a scalable and simple technique that shows strong performances in standard benchmarks for regression and classification.", "questions": "None, aside from the minor points presented in the weaknesses section", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. I found the paper very interesting and easy to read. Uncertainty estimation is an important and active research area in deep learning.\n\n1. To the best of my knowledge, the idea is novel (although not groundbreaking)\n\n1. The method is scalable, simple to implement in standard architectures, and achieves very competitive performances (especially considering its simplicity)\n\n1. While being Bayesian only on the last neural network layer the method is in principle not as powerful as other techniques, as the authors rightly claim simpler methods that are easy to implement are what is being more commonly used in practice (e.g. Bayesian dropout, stochastic weight averaging)\n\n1. The appendix is extensive and addresses all the details I felt were missing in the main paper", "weaknesses": "I did not identify any major weaknesses, only some points for improvement\n\n1. To increase the impact of the paper you need to make sure that people that are not too familiar with VI are able to easily implement the paper. This means:\n    1. Make the code publicly available, especially to show how to best implement the \"mixed parameterization\" discussed in appendix D \n    1. set good default hyperparameters\n\n1. It would be useful to draw the graphical models of the models presented in Section 2, to help the reader visualize the random variables in play and their (hyper)priors/parameters\n\n1. In case you need more space in the paper, I would move to the appendix some of the details on the generative classification model, especially considering the poorer performances.", "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 (spotlight)