Sampling-based inference for large linear models, with application to linearised Laplace
Authors: Javier Antoran
Shreyas Padhy
Riccardo Barbano
Eric Nalisnick
David Janz
José Miguel Hernández-Lobato
Venue: ICLR 2023 poster (2023)
Keywords: Laplace
linearised Laplace
Bayesian neural network
Bayesian linear regression
uncertainty estimation
Bayesian deep learning
EM
large scale regression
sample then optimise
evidence framework

Abstract: Large-scale linear models are ubiquitous throughout machine learning, with contemporary application as surrogate models for neural network uncertainty quantification; that is, the linearised Laplace method. Alas, the computational cost associated with Bayesian linear models constrains this method's application to small networks, small output spaces and small datasets. We address this limitation by introducing a scalable sample-based Bayesian inference method for conjugate Gaussian multi-output linear models, together with a matching method for hyperparameter (regularisation) selection. Furthermore, we use a classic feature normalisation method (the g-prior) to resolve a previously highlighted pathology of the linearised Laplace method. Together, these contributions allow us to perform linearised neural network inference with ResNet-18 on CIFAR100 (11M parameters, 100 output dimensions × 50k datapoints) and with a U-Net on a high-resolution tomographic reconstruction task (2M parameters, 251k output dimensions).

Published conference paper at ICLR 2023
SAMPLING-BASED INFERENCE FOR LARGE LINEAR MODELS, WITH APPLICATION TO LINEARISED LAPLACE
Javier Antorán"
Shreyas Padhy University of Cambridge
Riecardo Barbano
University of Cambridge
University College London
Eric Nalismick
David Janz University of Alherta
José Miguel Hernandez-Lobato University of Cambridge
University of Amsterdam
ABSTRACT
Large-scale linear models are ubiquitous throughout machine learning, with contem- porary application surrogate models for neural network uncertainty quantifica- tion; that is, the linearised Laplace method. Alas, the computational cost associated with Bayesian linear models constrains this method's 'ss application to small networks, small output spaces and small datasets. We address this limitation by introducing a scalable sample-based Bayesian inference method for conjugate Gaussian multi- output linear models, together with matching method for hyperparameter (regu- larisation strength) selection. Furthermore, we classic feature normalisation method, the g g-prior, to resolve a previously highlighted pathology of the linearised Laplace method. Together, these contributions allow us to perform linearised neural network inference with ResNet-18 18 CIFARIC 100 (1 parameters, 100 output
dimensions x 50k datapoints) and with a U-Net on a high-resolution tomographic reconstruction task (2M parameters, 251k output dimensions). An extended version of this work is available at arxiv wwx.v..rggpof..fooooooooooo 2210 04994 pddf.
INTRODUCTION
The linearised Laplace method, originally introduced by Mackay (1992), has received renewed interest in the context of uncertainty quantification for modern neural networks (NN) (Khan et al., 2019; Immer al., 2021b; Daxberger et al., 2021a). The method constructs surrogate Gaussian linear model for the NN predictions, and uses the error bars of that linear model estimates of the NN's uncertainty. However, the resulting linear model is very large; the design matrix is sized number of parameters by number of datapoints times number of output classes. Thus, both the primal (weight space) and dual (observation space) formulations the linear model are intractable. This restricts the method to small network or small data settings. Moreover, the method sensitive to the choice of regularisation strength for the linear model (Immer al., 2021a; Antorán et al.. 2022c). Motivated by linearised Laplace, we study inference and hyperparameter selection in large linear models.
To scale inference and hyperparameter selection in Gaussian linear regression, we introduce sample- based Expectation Maximisation (EM) algorithm. It interleaves E-steps, where we infer the model's posterior distribution over parameters given some choice of hyperparameters, and M-steps, where the hyperparameters are improved given the current posterior. Our contributions here are two-fold:
We enable posterior sampling for large-scale conjugate Gaussian-linear models with a novel sample-then-opiimize objective, which we use to approximate the E-step.
We introduce a method for hyperparameter selection that requires only access to posterior samples, and not the full posterior distribution. This forms our M-step.
Combined, these allow to perform inference and hyperparameter selection by solving a series of quadratic optimisation problems using iterative optimisation, and thus avoiding an explicit cubic cost in any of the problem's 's properties. Our method readily extends to non-conjugate settings, such as classification problems, through the use of the Laplace approximation. In the context of linearised
"qqual contribution. Correspondence 6a66 a.cm...cmm..mmmmmmmmmmmmmmm. ac. and sp2058 8cam. www.aaa.uu

---

Published conference paper at ICLR 2023
NNs, our approach also differs from previous work in that it avoids instantiating the full NN Jacobian matrix, an operation requiring as many backward passes as output dimensions in the network.
We demonstrate the strength of our inference technique in the context of the linearised Laplace procedure for image classification parameter ResNet-18. 18. We also consider high-resolution (25lk pixel) tomographic reconstruction (regression) task with a 2M parameter U-Net. In tackling these, we encounter a pathology in the M-step of the procedure first highlighted by Antorán al. (2022c): the standard objective therein ill-defined when the NN contains normalisation layers. Rather than using the solution proposed in Antorán et al. (2022c), which introduces more hyperparameters, we show that standard feature- normalisation method, the g-prior (Zellner. 1986; Minka, 2000), resolves this pathology. For the tomographic reconstruction task, the regression problem requires a dual-form formulation our E-step; interestingly, we show that this is equivalent to optimisation viewpoint on Matheron's 's rule (Journel & Huijbregts, 1978; Wilson al., 2020), connection we believe to be novel.
CIFARIO0 (100 classes 50k datapoints) using 11M
2 CONJUGATE GAUSSIAN REGRESSION AND THE EM ALGORITHM
We study Bayesian conjugate Gaussian linear regression with multidimensional outputs, where we observe inputs 2i,..., In E IRd4 and corresponding outputs ;,,...; E R.". We model these as
2/ o(x;)0 + T/s
(1)
where d: Rd - R x Rd is known embedding function. The parameters 0 are assumed sampled from,N (0, 1) with an unknown precision matrix A E ge'xe" and for each i S yΙ (0, BF²) additive noise vectors with precision matrices B3 E IRm relating the m output dimensions.
Our goal is to infer the posterior distribution for the parameters 0 given our observations, under the setting of A of the form A = al for > 0 most likely to have generated the observed data. For this, we use the iterative procedure of Mackay (1992). which alternates computing the posterior for denoted 11, for a given choice A, and updating A, until the pair (A, II) converge to a locally optimal setting. This corresponds to an EMM algorithm (Bishop, 2006).
Henceforth, we will use the following stacked notation: we write Y E Imm for the concate- nation of V1, Vns B E mm for block diagonal matrix with blocks B B, and (o(X,)" :... : (xx,) E Ramxd for the embedded design matrix. We write M = $ B&, Additionally, for a vector and PSD matrix G' of compatible dimensions, ||l|| = x Gu.
With that, the aforementioned EM algorithm starts with some initial E R2 xd and iterates:
(E step) Given A, the posterior for 0, denoted II, is computed exactly as
I= ( where H + and i θ= -
(M step) We lower bound the log-probability density the observed data, i.e. the evidence, for the model with posterior II and precision A' as (derivation in Appendix B.2)
log (Y; A') N 1 10 1 log det(I + A'
M(A),,
(3)
for C independent of A'. We choose an A that improves this lower bound.
Limited scalability The above inference and hyperparameter selection procedure for II and A is futile when both and nm are large. The E-step requires the inversion of d' x d' matrix and the M-step evaluating its log-determinant. both cubic operations in d'. These may be rewritten to instead yield cubic dependence nm (as in Section 3.3), but under assumptions, that too is not computationally tractable. Instead, we now pursue stochastic approximation to this pM-prrdduuee.
3 EVIDENCE MAXIMISATION USING STOCHASTIC APPROXIMATION
We now present our main contribution, a stochastic approximation (Nielsen, 2000) to the iterative algorithm presented in the previous section. Our M M-step requires only access to samples from II. We introduce method approximate posterior samples through stochastic optimisation for the E-step.

---

Published conference paper at ICLR 2023
3.1 HYPERPARAMETER SELECTION USING POSTERIOR SAMPLES (M-STEP)
For now, assume that we have an efficient method of obtaining samples (i...., Sk ~ II at each step. where I1 is zero-mean version of the posterior I1, and access to 6, the mean of II. Evaluating the first order optimality condition for M (see Appendix B.3) yields that the optimal choice of satisfies
|@1 Tr [r %
(4)
where the quantity 7 is the effective dimension of the regression problem. It can be interpreted as the number of directions in which the weights 0 are strongly determined by the data. Setting A = al for >/ | yields contraction step converging towards the optimum of M (Mackay, 1992).
Computing y directly requires the inversion of H, cubic operation. We instead rewrite y as an expectation with respect to 11 using Hutchinson (1990)'s 's trick, and approximate it using samples as
7 - Tr (1 M) TT( MH #} wKç M(1] R ビ CJ 4T BBC, =
(5)
We then select 3/11012- We have thus avoided the explicit cubic cost of computing the log- determinant in the expression for M (given in (3)) or inverting H. Due to the block structure of B, y may be computed in order vector-matrix products.
3.2 SAMPLING FROM THE LINEAR MODEL'S POSTERIOR USING SGD (E-STEP)
Now we turn to sampling from II" = N0 H It is known (also, shown in Appendix C.1) that for E Rnm the concatenation Ei, En with E; ~ (0, B, and 6 ~ NO A the minimiser of the following loss is random variable ( with distribution II"
((=) 1 | a = 09113-
(6)
This is called the "sample- thepppleette--pppp optimise" method (Papandreou & Yuille, 2010). We may thus obtain a posterior sample by optimising this quadratic loss for given sample pair (E, 00). Examining L:
The first term data dependent. It corresponds to the scaled squared errrr in fitting & as a linear combination of 4. Its gradient requires stochastic approximation for large datasets.
The second term, a regulariser centred at 60, does not depend on the data. gradient can thus be computed exactly at every optimisation step.
Predicting & i.e. random noise, from features 4 is hard. Due to this, the variance of mini-batch estimate of the gradient of || Eli may he large. Instead, for & and ( defined before, we propose the following alternative loss, equal to L up to an additive constant independent of 2:
(7) The mini-batch gradients of L' and L are equal in expectation (see Appendix C.1). However, in L', the randomness from the noise samples & and the prior sample 1 both feature within the regularisation term-uhh the gradient which can be computed exactly-rather than the data-dependent term. This may lead to lower minibatch gradient variance. To see this, consider the variance of the single-datapoint stochastic gradient estimators for both objectives' data dependent terms. At E Rd for datapoint index j ~ Unif({1,.... ..n......,)) these are
E'(z)= 1 2 112 + +|l= 13 with == - A BEE
no(X,) (o(X))= (xX))) Eạ) and D' - ===n((X) (=xxxxX)(xXx))) o(X))=
(8)
for L and L', respectively. Direct calculation, presented in Appendix C.2, shows that 1 [Var - Varğ'l Var(*"BE) 2Cov(""BBe +TBB)= A.
(9)
Note that both Varğ and Varğ' are d' d' matrices. We impose an order on these by considering their traces: we prefer the new gradient estimator g' if the sum of its per-dimension variances is lower than that of ÿ: that is if Tr A O. We analyse two key settings:
initialisation. taking 8 (or any other initialisation independent of E).
A = Tr / 138 Tr [+* B0ED"" (BB)) Tr (M) > 0,
(10)
where we used that E[EE-] B and since E is zero mean and independent of 8o, we have B R Thus, the new objective 1.' is always preferred initialisation.

---

Published conference paper at ICLR 2023
convergence, that is, when Ç. more involved calculation presented Appendix C.3 shows that L' is preferred if
20ry > Tr M.
(11)
This satisfied if a is large relative to the eigenvalues of M (see Appendix C.4), that is, when the parameters are not strongly determined by the data relative to the prior.
When 1'' is preferred both at initialisation and at convergence, we expect it to have lower variance for most minibatches throughout training. Even if the proposed objective I/ is not preferred convergence, may still be preferred for most the optimisation, before the noise is fit well enough.
3.3 DUAL FORM OF E-STEP: MATHERON'S RULE OPTIMISATION
The minimiser of both and is (A00 x BE). The dual (kernelised) form of this is
<= 0 + A 1--4""00
" (B 4
(s 400))
(12)
which is known in the literature Matheron's rule (Wilson al., 2020, Appendix D). Evaluating (12) requires solving mm-dimensional quadratic optimisation problem, which may be preferable the primal problem for mn < d"; however, this dual form cannot be minibatched over observations. For small can solve this optimisation problem using iterative full-bate full-batch quadratic optimisation algorithms (e.g. conjugate gradients), significantly accelerating our sample-based EM iteration.
4 NN UNCERTAINTY QUANTIFICATION AS LINEAR MODEL INFERENCE
Consider the problem of m-output prediction. Suppose that we have trained neural network of the form f: Fd Rd RT, obtaining weights w E 12" using loss of the form
[(f(w,-)) Σ.= Ev, fw, x.)) R(w)
(13)
where & a data term negative log-likelihoon and linearised Laplace to quantity uncertainty in the network predictions f(w, We then present the g e-prior, a feature normalisation that resolves certain pathology in the linearised Laplace method when the network contains normalisation layers.
regulariser. now show how to use
4.1 THE LINEARISED LAPLACE METHOD
The linearised Laplace method consists of two consecutive approximations, the latter of which is necessary only if is non-quadratic (that is, if the likelihood is non- :ausssai))):
We take a first-order Taylor expansion of around vi, yielding the surrogate model
h(0, x) f( x) (0z) - () for ((z) (w,
(14)
This is an affine model in the features 6(x) given by the network Jacobian at x.
We approximate the loss of the linear model L(h(6.-) with a quadratic, and treat it as a negative log- den-eiyyty for the parameters 0, yielding a Gaussian posterior of the form
NOO (V -V -))
where
Ó E argming ((h(0,-)).
(15)
Direct calculation shows that (V)C)(((,)) -)) (A 4 H1, for V R(()) and B a block diagonal matrix with blocks B, V3, E(v:. evaluated at O 1(0.x;).
We have thus recovered a conjugate Gaussian multi-output linear model. We treat A as a learnable parameter thereafter IIn practice, we depart from the above procedure in two ways:
We use the neural network output fw, the predictive mean, rather than the surrogate model mean (00 -). Nonetheless, still need to compute 0 to within the M-step of the EM procedure. To do this, we minimise C(h(0, -)) over 0 E R
'The EM procedure from Section 2 for the conjugate Gaussian-lines model, where carries guarantces non-decreasing model evidence, and thus convergence local optimum. These guarantees do not hold for non-conjugate likelihood functions, e.g., the wwttaa-cateeoriice where the Laplace approximation necessary.

---

Published conference paper at ICLR 2023
We compute the loss curvature B, at predictions y; = f(w, place h(d, zz),, since the latter would change each time the regulariser A is updated, requiring expensive re-evaluation.
Both of these departures are recommended within the literature (Antorán et al., 2022c).
4.2
ON THE COMPUTATIONAL. ADVANTAGE OF SAMPLE-BASED PREDICTIONS
The lincarised Laplace predictive posterior at an input is [V (1 2), ((x) H=(=)""))))")) Even given H evaluating this naively requires instantiating ((2), cost of vector-Jacobian products (i.e. backward passes). This is prohibitive for large m. However, expectations of any function
r R + IR under the predictive posterior can be approximated using only samples from IT as
Bir] ~ > r()))) for 3 f(w, x) + ((x){; with SI..... ~ II,
(16)
requiring only k Jacobian-vector products. In practice, we find l much smaller than m suffices.
4.3 FEATURE EMBEDDING NORMALISATION: THE DIAGONAL G-PRIOR
Due to symmetries and indeterminacies neural networks, the embedding function x(-) = V used in linearised Laplace method yields features with arbitrary scales across the d' weight dimensions. Consequently, the dimensions of the embeddings may have an (arbitrarily) unequal weight under isotropic prior; that is, considering (()) for 6 ~ (0,
There are two natural solutions: either normalise the features by their (empirical) second moment, resulting in the normalised embedding function ø' given by
o' (x) (()) diag(s) for E Rd given by S; s; - B+) [+'BB],
(17)
likewise scale the prior, setting A (iiag((s" () The latter formulation is diagonal version what is known the literature as g-prior (Zellner, 1986) or invariant prior (Minka, 2000).
The g-prior may, in general. improve the conditioning of the linear system. Furthermore. when the linearised network contains normalisation layers, such as batchnorm (that is, most modern networks), the g-prior is essential. Antorán al. (2022c) show that normalisation layers lead to indeterminacies NN Jacobians, that in turn lead to an ill-defined model evidence objective. They propose learning separate regularisation parameters for each normalised layer of the network. While fixing the pathology, this increases the complexity of model evidence optimisation. As show in Appendix E, the g-prior cancels these indeterminacies, allowing the use a single regularisation parameter.
5 DEMONSTRATION SAMPLE-BASER LINEARISED LAPLACE INFERENCE
We demonstrate our linear model inference and hyperparameter selection approach on the problem estimating the uncertainty in NN predictions with the linearised Laplace method. First, in Section 5.1, we perform an ablation analysis on the different components of algorithm using small LeNet-style CNNs trained on MNIST. In this setting, full-covariance Laplace inference (that is, exact linear model inference) is tractable, allowing to evaluate the quality of our approximations. We then
Algorithm 1: Sampling- based linearised Laplace hyperparameter learning and inference
Inputs: initial 0; A A E N. number of samples for stochastic EM and prediction. respectively.
Compute g-prior scaling vector as in (17)
Sample random regularisers 0f.....,0" per (7)
while a has converged do
Find posterior mode 0 by optimising linear model loss L(h(0,)) given in (13) Draw posterior samples ....G. &u by optimising objective L' with 07 ,f.....77 86 Estimate effective dimension 7 per (5), using samples {1...6.
Update prior precision a + -/
Sample random regularisers 0f' ..... 0g' using optimised a
Draw corresponding posterior samples ci...., So using loss /!'
Output: posterior samples Ci.....

---

Published conference paper at ICLR 2023
peior functiotas
samppl optimizaaiin
poetector Sunctioras
bedure EM a 0.001
EM
24
2-
0.31
10-
0
Figure 1: Illustration of our procedure for fully connected NN the toy dataset of Antorán al. (2020). Top: prior function samples present large std-dev. (left). When these samples are optimised (middle shows a 2D slice of weight space), the resulting predictive errorbars are larger than the marginal target variance (right). Bottom: after EM, the std-dev. of prior functions roughly matches that of the targets (left), the overlap between prior and posterior is maximised, leading to shorter sample trajectories (center), and the predictive errorbars are qualitatively more appealing (right).
demonstrate our method at scale on CIFAR100 classification with ResNet18 (Section 5.2) and the dual (kernelised) formulation of our method on tomographic image reconstruction using a U-Net (Section 5.3). We look at both marginal and joint uncertainty calibration and at computational cost.
For all experiments, our method avoids storing covariance matrices H computing their log- determinants, or instantiating Jacobian matrices Φ(z). all of which have hindered previous linearised Laplace implementations. We interact with NN Jacobians only through Jacobian-vector and vector- Jacobian products, which have the same asymptotic computational and memory costs NN forward-pass (Novak al., 2022). Unless otherwise specified, we use the diagonal g-prior and a scalar regularisation parameter. Algorithm summarises our method, Figure shows an illustrative example, and full algorithmic detail is in Appendix F. An implementation of our method in JAX can be found here. Additional experimental results are provided in Appendix and Appendix I.
5.1 ABLATION STUDY: LENET ON MNIST
We first evaluate our approach on MNIST m=10 class image classification, where exact linearised Laplace inference is tractable. The training set consists of n=60k observations and we employ 3 LeNet-style CNNs of increasing size: "LeNetSmall" 44344, "L.eNet" (d' (d'=29226) and "LeNet- Big" (d'=46024). The latter is the largest model for which we can store the covariance matrix on an A100 GPU, We draw samples and estimate posterior modes using SGD with Nesterov momentum (full details in Appendix G). We use 5 seeds for each experiment, and report the mean and std. error.
Comparing sampling objectives We first compare the proposed objective L' with the one standard the literature L., using LeNet. The results are shown Figure 2. We draw exact samples: (C) jıs through matrix inversion and assess sample fidelity in terms of normalised squared distance to these exact samples IK- 11/11< 113. All runs share prior precision of a 5.5 obtained with EM iteration. The effective dimension is 7 Rs 1300. Noting that the g-prior feature normalisation results in T = d', we can see that condition (11) is not satisfied (2 x 5.5 x 1300 < 29k).. Despite this, accurate samples even when using 16-times smaller batch
<<<<'11//1K""
the proposed objective converges size (left plot). The right side plots relate sample error to categorical symmetrised -LL (sym. KL) and logit Wasserstein-2 (W2) distance between the sampled and exact lin. Laplace predictive distributions on the test-set. Both objectives' prediction errors stop decreasing below sample error of 20.5 but, nevertheless, the proposed loss L/' reaches lower a prediction error.
Fidelity of sampling inference We compare our method's 's uncertainty using 64 samples against approximate methods based on the NN weight point-estimate (MAP), a diagonal covariance, and a KFAC estimate of the covariance (Martens & Grosse, 2015; Ritter et al., 2018) implemented with the Laplace library, in terms of similarity to the coll-covariance lin. Laplace predictive posterior. As standard, we compute categorical predictive distributions with the probit approximation (Daxberger al., 20) All methods use the same layerwise prior precision obtained with steps of full- covariance EM iteration. For all three LeNet sizes, the sampled approximation presents the lowest categorical sym. KL and logit W2 distance to the exact lin. Laplace pred. posterior (Figure 3, LHS). The fidelity of competing approximations degrades with model size but that of sampling increases.

---

Published conference paper at ICLR 2023
sample error norm vs batch size
pred. err. vs sample err.
old
10:2
batch size
5X1 KL
10P
10:3 3
W2
10"4
le44
2000
10-3
4000
2000
10
115
11
train steps
Figure 2: Left: optimisation traces for new and existing sampling losses averaged across 16 samples and 5 seeds. Right: for watch-size=1000 1000, traces of sample-error distance to exact linear predictions.
difference with full lin. Laplace predictive
isotropic precision
layerwise precision
311
A sempling
10-3
prior prec.
prier prus.
KL
W2
107
10-1
102
LaNoiSmall LoNad
LaNaBBig
LeNeSwulll
LeNes
LoNetiig
FM step
EM
# 1-1444 89-29126
40024
4' 14634
2222026
44-46004
exect
KTAC
Iayeno
laper
sampling
KBAC
ding
MAP
sampling
Laysr
Jayse
Figure 3: Left: similarity to exact lin. Laplace predictions on MNIST sest-set for different approximate methods applied to NNs of increasing size. Centre right: comparison of EM convergence for single hyperparameter across approximations. Right: layerwise convergence for exact and sampling methods.
compare our sampling EM iteration with 16 samples to full-covariance EM on LeNet without the g-prior. Figure 3, right, shows that for a single precision hyperparameter, both approaches converge in about 3 steps the same value. In this setting, the diagonal covariance approximation diverges, and KFAC converges to biased solution. We also consider learning layer-wise prior precisions by extending the M-step update from Section 3 to diagonal but non-isotropic prior precision matrices (see Appendix B.4), Here, neither the full covariance nor sampling methods converge within 15 EM steps. The precisions for all but the final layer grow, revealing a pathology of this prior parametrisation: only the final layer's Jacobian, i.e. the final layer activations, needed to accurately predict the targets; other features are pruned away.
Accuracy of sampling hyperparameter selection
5.2 RESNET18 8 ON CIFAR 100
We showcase stability and performance our approach by applying it to CIFAR 100 m= way image classification. The training set consists of n=50k observations, and we employ a ResNet-18 model with d' ss 11M parameters. We perform optimisation using SGD with Nesterov momentum linear learning rate decay schedule. Unless specified otherwise, we run 8 steps EM with 6 samples to select a. We then optimise 64 samples to be used for prediction. We run each experiment with 5 different seeds reporting mean and std. error. Full experimental details in Appendix G.
Stability and cost of sampling algorithm Figure 4 shows that our sample-based EM converges 6 steps, even when using a single sample. At convergence, & 10 and fy Ps 700, so 2 x 700 10
1Ex
pror prec.
eff. dim 4
avg
103
samples
0.0
EM steps
EM steps
EM steps
Figure 4: Left: prior precision optimisation traces for ResNet18 8 on CIFAR100 varying samples. Middle: same for the eff. dim. Right: average sample and posterior norm throughout successive EM steps' SGD runs while varying n. samples. Note that traces almost perfectly overlap.

---

Published a conference paper at ICLR 2023
MAP
Ensemble (5)
KBAC
Sampling
test LL
marginal I.I.
-1.40
-030
-.12 0:01
-1077 0.00
-13.97 4: 0.01
-6.56 a: 0.01
4.92 = 0.04
-5.14
KFAC
joint LL
-27.88 + 0.08
-14.17
0.00
-10.3 +.183++0.12
-10.77 + 0.09
41.88 0.03
22.20
0.04
-1.02
0.22
-18.04
1.118
-35.89
0.02
-3.07 0:09
-2..00
-20.75
0.28
Phobit
corruption severity
Figure 5: Performance under distribu- tion shift for ResNet18 and CIFARI00.
Table I: Comparison of methods' marginal and joint pre- diction performance for ResNet18 on CIFAR 100.
1.4 x 10 > 1.1 x 10 Thus, (11) is satisfied and 1'' is preferred. We use 50 epochs of optimisation for the posterior mode and 20 for sampling. When using 2 samples, the cost of one EM step with our method is 45 minutes on an A100 GPU; for the KFAC approximation. this takes 20 minutes.
Evaluating performance in the face of distribution shift We employ the standard benchmark for evaluating methods' test Log-Likelihodd (LLL) the increasingly corrupted data sets of Hendrycks Gimpel (2017); Ovadia et al. (2019). We compare the predictions made with our approach to those from 5-element deep ensembles, arguably the strongest baseline for uncertainty quantification in deep learning (Lakshminarayanan et al., 2017: Ashukha et al., 2020), with point-estimated predictions (MAP), and with a KFAC upproximation of the lin. Laplace covariance (Ritter et al., 2018). For the latter, constructing full Jacobian matrices for every test point is computationally intractable, so we 64 samples for prediction, as suggested in Section 4.2. The KFAC covariance structure leads fast log-determinant computation, allowing us to learn layer-wise prior precisions (following Immer et al., la) for this baseline using 10 steps of non-sampled EM. For both lin. Laplace methods, we use the standard probit approximation to the categorical predictive (Daxberger al., 2021b). Figure 5 shows that for in-distribution inputs, ensembles performs best and KFAC overestimates uncertainty, degrading LL even relative to roint-estimated MAP predictions. Conversely, our method improves LL. For sufficiently corrupted data, our approach outperforms ensembles, also edging out KFAC, which fares well here due to its consistent overestimation of uncertainty.
Joint predictions Joint predictions are essential for sequential decision making, but are often ignored in the context of NN uncertainty quantification (Janz et al., 2019). To address this, we replicate the "dyadic sampling" experiment proposed by Osband et al. (2022). We group our test-set into sets of k data points and then uniformly re-sample the points in each set until sets contain 7 points. We then evaluate the LL of each set jointly. Since each only contains distinct points, a predictor that models seff-covariances perfectly should obtain an LL value at least as large as its marginal LL for all values of k. We use 1=10(x 1) and repeat the experiment for 10 test-set shuffles. Our setup remains the same as above but we Monte Carlo marginalisation instead of probit, since the latter discards covariance information. Table shows that ensembles make calibrated predictions marginally but their joint predictions are poor, an observation also made by Osband et al. (2021). Our approach is competitive for all x, performing best in the challenging large cases.
5.3 TOMOGRAPHIC RECONSTRUCTION
To demonstrate our approach in dual (kernliised) form, we replicate the setting of Barbano et al. (2022a:b) and Antorán et al. (2022b), where linearised Laplace is used estimate uncertainty for a tomographic reconstruction outputted by a U-Net autoencoder. We provide an overview of the problem in Appendix G.3, but refer to Antorán et al. (2022b) for full detail. Whereas the authors use single EM step learn hyperparameters, otherwise specified, we use 16 samples for stochastic EM, and 1024 for prediction.
use our sample-basse based variant and run 5 steps. Unless
We test on the real-measured JCT dataset of 251k pixel seans of a single walnut released by Der Sarkissian et al. (2019b). We train d' M parameter U-Nets on the m=7680 dimensional observation used by Antorán al. (2022b) and a twice as large setting m=15360. Here, the U-Net's imput is clamped to a constant (n = 1), and its parameters are optimised to output the reconstructed image. As result, we do not need mini-batchhie batching and can draw samples using Matheron's 's rule (12). We solve the linear system contained therein using conjugate gradient (CG) iteration implemented

---

Published conference paper at ICLR 2023
all. din
0000
134
strps
steps
steps
patchesin
Figure 6: Left 3 plots: traces prior precision, eff. dim., and marginal test LL vs EM steps for the tomographic reconstruction task with 7680. Right: joint test LL for varying image patch sizes.
Table 2: Tomographic reconstruction: test LL and wall-clock times (A100 GPU) for both data sizes.
7680
15366)
wall-clock (min.)
wall-clock (min.)
Method
marginal
(10 10)
purams optim.
prediction
marginal
(10 10)
params, optim. prediction
0.028
2.762
lin.-UNer
2.214
2.601
1260'
196'
sampl--lin-UNN
2.341
2.972
with CPyTorch (Gardner et al., 2018). We accelerate CG with randomised SVD preconditioner rank 400 (alg. 5.6 in Halko et al., 2011). See Appendix G for full experimental details.
Stability and cost of sampling algorithm Figure shows that cample-based EM iteration converges within 4 steps using few as samples. Table shows the time taken to perform 5 sample-based EM steps for both the m=7680 and m= 15365 settings; avoiding explicit estimation of the covariance log-determinant provides with two order of magnitude speedup relative to Antorán et al. (2022c) for hyperparameter learning. By avoiding covariance inversion, we obtain order of magnitude speedup for prediction. Furthermore, while scaling to double the observations m=15330 intractable with the previous method, our sampling method requires only 25% increase in computation time.
Predictive performance Figure 7 shows, qualitatively, that the marginal standard deviation assigned each pixel by our method aligns with the pixelwise error in the U-Net reconstruction in a fine- grained manner. By contrast, MC dropout (MCDO), the most common baseline for NN uncertainty estimation in tomographic reconstruction (Laves et al., 2020; Tölle et al., 2021), spreads uncertainty uniformly across large sections of the image. Tahle 2 shows that the pixelwise L.I. obtained with our method exceeds that obtained by Antorán et al. (2022b), potentially due to us optimising the prior precision convergence while the previous work could only afford single EM step. The rightmost plot Figure displays joint test LL, evaluated on patches of neighbouring pixels. Our method performs best. MCDO's predictions are poor marginally. They improve when considering covariances, although remaining worse than lin. Laplace.
Ix*
f(w3
sid-dev
1K M(J0)
sid-dev
sampp lin. UNir
MCDO-UNN
13
308
Figure 7: Original 501 x501 pixel walnut image and reconstruction error for a 7=7660 dimensional observation, along with pixel-wise std-dev obtained with sampling lin. Laplace and MCDO.
CONCLUSION
Our work introduced sample-based approximation to inference and hyperparameter selection in Gaussian linear multi-output models. This allowed scale the linearised Laplace method ResNet-18 on CIFAR100, where it was computationally intractable with existing methods. We also demonstrated the strength of the approach where it decreases the cost of hyperparameter selection by two orders of magnitude. The uncertainty estimates obtained through our method are well-calibrated not just marginally, but also jointly across predictions. Thus, our work may be of interest in the fields of active and reinforcement learning, where joint predictions of importance, and computation of posterior samples is often needed.
high resolution tomographic reconstruction task.

---

Published conference paper at ICLR 2023
REPRODUCIBILITY STATEMENT
In order to aid the reproduction of our results, we provide a high-level overview of our procedure in in two major experiments in Appendix F.
algorithm 1 and the fully detailed algorithms Appendix G provides full experimental details for all datasets and models used our experiments. Our code is available in repository at this link.
ACKNOWLEDGEMENTS
The authors would like to thank Alex Terenin and Marine Schimel for helpful discussions. JA acknowledges support from Microsoft Research, through its PhD Scholarship Programme, and from the EPSRC. SP acknowledges support from the Harding Distinguished Postgraduate Schol- ars Programme Leverage Scheme. JMHL acknowledges support from Turing Al Fellowship under grant EP/V023756/1. This work has been performed using resources provided by the Cambridge Tier-2 system operated by the University of Cambridge Research Computing Service thtp://www.hpc.comm www.hpccccm.c..kk cak.ac.uk funded by an EPSRC Tier-2 capital grant. This work was also supported with Cloud TPUs from Google's S TPU Research Cloud (TRC).
REFERENCES
Vincent Adam, Paul E. Chang, Mohammad Emtiyaz Khan, and Arno Solin. Dual parameterization of sparse variational gaussian processes. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Perey Liang. and Jennifer Wortman Vaughan (eds.), Advances Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NewrlPS 2021, December 6-14 14, 2021, virtual, pp. 11474-11486, 2021.
Javier Antorán, James Allingham, and José Miguel Hernánezz-Lobato. Depth uncertainty in neural networks. H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin (eds.). Advances Neural Information Processing Systems, volume 33, pp. 106344 Curran Associates, Inc., 2020.
Javier Antorán, James Urquhart Allingham, David Janz, Erik Daxberger, Eric Nalisnick, and José Miguel Hernández-Lobaad Linearised laplace inference in networks with normalisation layers and the neural g-prior. In Fourth Symposium on Advances in Approximate Bayesian Infer- ence, 2022a.
Javier Antorán, Riccardo Barbano, Johannes Leuschner, José Miguel Hernández-Lobato, and Bangti Jin. probabilistic deep image prior for computational tomography. CoRR, abs/2203.00479, 2022b. doi: wtwp://aawxi..2..0044..
Javier Antorán, David Janz, James Urquhart Allingham, Erik A. Daxberger, Riccardo Barbano, Eric T. Nalisnick, and José Miguel Hemández--Loalle Adapting the linearised laplace model evidence for modern deep learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song. Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.). International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 Proceedings of Machine Learning Research, pp. 196-8211 PMLR, 2022c.
Jordan T. Ash, Cyril Zhang, Surbhi Goel, Akshay Krishnamurthy, and Sham M. Kakade. Anti- concentrated confidence bonuses for scalable exploration. In International Conference on Learning Representattons, 2022.
Arsenii Ashukha, Alexander Lyzhov, Dmitry Molchanov, and Dmitry P. Vetrov. Pitfalls in-domain uncertainty estimation and ensembling in deep learning. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net.
2020.
Francis R. Bach. Adaptivity of averaged stochastic gradient descent to local strong convexity for logistic regression. J. Mach. Learn. Res., 15(1))595-627, 2014. doi: 10.555/267445.222744444
10

---

Published conference paper at ICLR 2023
Daniel Otero Baguer, Johannes Leuschner, and Maximilian Schmidt. Computed tomography recon- struction using deep image prior and learned reconstruction methods. Inverse Problems, 36(9): 094004, 2020.
M. Baragatti and D. Pommeret. A study of variable selection using 8-prior distribution with ridge parameter. Computational Statistics & Data Analysis, 56(6):1220-1934 2012. ISSN 0167-9473. doi: https://doi. 10. 101 (016//..sst 2011..
Riccardo Barbano, Johannes Leuschner, Maximilian Schmidt, Alexander Denker, Andreas Haupt- mann, Peter MaaB, and Bangti Jin. Is deep image prior in need of good education? arXiv preprint arXiv:2111. 19966, 2021.
Riccardo Barbano, Javier Antorán, José Miguel Srenändez-Lobatee and Bangti Jin. A probabilistic deep image prior over image space. Fourth Symposium on Advances Approximate Bayesian Inference, 2022a.
Riccardo Barbano, Johannes Leuschner, Javier Antorán, José Miguel Hernández- Lobato, and Bangti Jin. Bayesian experimental design for computed tomography with the linearised deep image prior. ICML Workshop Workshop on Adaptive Experimental Design and Active Learning in the Real World, 2022b.
Riccardo Barbano, Javier Antorán, Johannes Leuschner, José Miguel Hernandez-Lobato, Željko Kereta, and Bangti Jin, Fast and painless image reconstruction in deep image prior subspaces. 2023.
Christopher M Bishop. Pattern recognition and machine learning. springer, 2006.
Daniel Sabanés Bové and Leonhard Held. Hyper-g priors for generalized linear models. Bayesian Analysis, 6(3):387 410. 2011. doi: 10.1214/11-B6115
Frik Daxherger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace effortellss Bayesian deep learning. In M. Ranzato, A. Beygelzimer, P.S. Liang, J. Vaughan, and Dauphin (eds.), Advances in Neural Information Processing Systems, volume 34, 6-14 Dec 2021a.
Erik Daxberger, Eric Nalisnick, James Allingham, Javier Antorán, and Jose Miguel Hernandez- Lobato. Bayesian deep learning via subnetwork inference. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 2510-2521. PMLR, 18-24 Jul 2021b.
Alexander G. de G. Matthews, Jiri Hron, Richard E. Turner, and Zoubin Ghahramani. Sample-then- optimize posterior sampling for bayesian linear models, In NeurlPS Workshop on Advances Approximate Bayesian Inference, 2017.
Henri Der Sarkissian, Felix Lucka, Maureen Eijnatten, Giulia Colacicco, Sophia Coban, and Joost Batenburg. Cone-beam x-ray et data collection designed for machine learning: Samples 1-8. 2019a.
Henri Der Sarkissian, Felix Lucka, Maureen van Eijnatten. Giulia Colacicco, Sophia Bethany Coban, and K. Joost Batenburg. Cone-Beam X-Ray CT Data Collection Designed for Machine Leaminn: Samples 1-8, 2019b.
Aymeric Dieuleveut, Nicolas Flammarion, and Francis R. Bach. Harder, better, faster, stronger convergence rates for least-squares regression. Mach. Learn. Res., 18:101:1-101:51, 2017.
Runa Eschenhagen, Erik Daxberger, Philipp Hennig, and Agustinus Kristiadi. Mixtures of laplace approximations for improved post-hoc uncertainty in deep learning. CoRR, abs/21 1111111001 1.03577, 2021.
Jacob R. Gardner, Geoff Pleiss, Kilian Q Weinberger, David Bindel, and Andrew Gordon Wilson. gaussian process inference with GPU acceleration. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 31: Anmual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, 3-8 December 2018, Montréal. Canada, pp. 7587-7597. 2018
GPyTorch: Blackbox

---

Published a conference paper at ICLR 2023
Stephen F. Gull. Bayesian Data Analysis: Straight-line fittinnn pp. 511-518. Springer Netherlands, Dordrecht, 1989. ISBN 978-94-015-78604 doi: 11100//778-94-015-7868 55.
Nathan Halko, Per-Gunnar Martinsson, and Joel A. Tropp. Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions. SIAM Rev., 53(2): 217-288, 2011. doi: 10.1 10.1137/090778800
Dan Hendrycks and Kevin Gimpel. A Baseline for Detecting Misclassified and Out-of-Distribution Examples in Neural Networks, In International Conference on Learning Representations (ICLR), 2017.
James Hensman, Nicoló Fusi, and Neil D. Lawrence. Gaussian processes for big data. In Ann E. Nicholson and Padhraic Smyth (eds.). Proceedings of the Twenty-Ninth Conference on Uncertainty in Artificial Intelligence, UAI 2013, Bellevue, WA, USA, August 11-15, 2013. AUAT Press, 2013.
Yehuda Hoffman and Erez Ribak. Constrained Realizations Gaussian Fields: A Simple Algorithm. The Astrophysical Journal Letters, 380:15,, Octoher 1991. doi: (0.1086/161660
M.F. Hutchinson. stochastic estimator of the trace the influence matrix for laplacian smoothing splines. Communicattooss in Statistics Simulation and Computation, 19(2):433-450, 1990. doi:
10. 10.1080/036109199008 12866
Alexander Immer, Matthias Bauer, Vincent Fortuin, Gunnar Rätsch, and Mohammad Emtiyaz Khan. Scalable marginal likelihood estimation for model selection in deep learning. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 4563 4573. PMLR, 2021a.
Alexander Immer, Maciej Korzepa, and Matthias Bauer. Improving predictions of bayesian neural nets via local linearization. In Arindam Banerjee and Kenji Fukumizu (eds.). The 24th International Conference on Artificial Intelligence and Statistics, AISTATS 2021, April 13-15, 2021, Virtual Event, volume 130 of Proceedings of Machine Learning Research, pp. 703-71 1. PMLR, 2021 b.
Alexander Immer, Tycho F.A. van der Ouderaa, Gunnar Ratsch, Vincent Fortuin, and Mark van der Wilk. Invariance learning in deep neural networks with differentiable laplace approximations. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho Advancce in Neural Information Processing Systems, 2022.
Pavel Izmailov, Sharad Vikram, Matthew D Hoffman, and Andrew Gordon Gordon Wilson. What are bayesian neural network posteriors really like? In Marina Meila and Tong Zhang (eds.), Proceed- ings of the 38th International Conference on Machine Learning. volume 139 of Proceedings of Machine Learning Research, pp. 4629 4640. PMLR, 18-24 24 Jul 2021.
Arthur Jacot, Clément Hongler, and Franck Gabriel. Neural tangent kernel: Convergence and generalization in neural networks. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa- Bianchi, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pp. 8580-8589, 2018
David Janz, Jiri Hron, Przemyslaw Mazur, Katja Hofmann, José Miguel Bernández-Lobate and Sebastian Tschiatschek. Successor uncertainties: Exploration and uncertainty in temporal differ- ence learning. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buec Buc, Emily B. Fox, and Roman Garnett (eds.). Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 4509-4518, 2019.
A. G. Journel and C. J. Huijbregts. Mining geostatistics / [by] A. G. Journel and Ch. J. Huijbregts. Academic Press London ; New York, 1978. ISBN 0123910501.
Mohammad Emtiyaz Khan, Alexander Immer, Ehsan Abedi, and Maciej Korzepa. Approximate inference turns deep networks into gaussian processes. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché- Emily B. Fox, and Roman Gamett (eds.), Advances in Neural Information Processing Systems 32; Annual Conference on Neural Information Processing
12

---

Published conference paper at ICLR 2023
Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC. Canada, pp. 2019.
Agustinus Kristiadi, Matthias Hein, and Philipp Hennig. Being bayesian, even just a bit, fixes overconfidence relu networks. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 5436- 0446..546.66 PMLR, 2020.
Frederik Kunstner, Philipp Hennig, and Lukas Balles, Limitations of the empirical fisher approxi- mation for natural gradient descent. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buc, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NewrlPS 2019, December 8- 14, 2019, Vancouver, BC, Canada, pp. 4158 4169, 2019.
Balaji Lakshminarayanan. Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Advances in Neural Information Processing Systems, pp. 6402-6413, 2017.
Max-Heinrich Laves, Malte Tölle, and Tobias Ortmaier. Uncertainty estimation medical image denoising with bayesian deep image prior. In Uncertainty for Safe Utilization of Machine Learning in Medical Imaging. and Graphs in Biomedical Image Analysis, pp. 81-96. Springer, 2020.
Neil D. Lawrence. Variational inference in probabilistic models. PhD thesis, University of Cambridge. 2000.
Jaehoon Lee, Lechao Xiao, Samuel S. Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer. Florence A'Alhhé-Be Buc, Emily Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 8570-8581, 2019.
Feng Liang, Rui Paulo, German Molina, Merlise A Clyde, and Jim Berger. Mixtures of g priors for bayesian variable selection. Journal of the American Statistical Association, (03448)):400423, 2008. doi:
David John Cameron Mackay. Bayesian Methods for Adaptive Models. PhD thesis, USA, 1992.
Wesley Maddox, Shuai Tang, Pablo G. Moreno, Andrew Gordon Wilson, and Andreas Damianou. Fast adaptation with linearized neural networks. Arindam Banerjee and Kenji Fukumizu (eds.), The 24th International Conference on Artificial Intelligence and Statistics, AISTATS 2021, April 13- 15, 2021, Virtual Event, volume 130 of Proceedings of Machine Learning Research, pp. 2737-2745. PMLR, 2021.
Wesley J. Maddox. Pavel Izmailov, Timur Garipov, Dmitry P. Vetrov, and Andrew Gordon Wilson. A simple baseline for bayesian uncertainty in deep learning. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence Alché-Buce Emily B. Fox, and Roman Garnett (eds.). Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 13132-13143, 2019.
Wesley J. Maddox, Gregory W. Benton, and Andrew Gordon Wilson. Rethinking parameter counting deep models: Effective dimensionality revisited. CoRR, abs/2003.02 139, 2020.
James Martens and Roger B. Grosse. Optimizing neural networks with kronecker-factoree factored approxi- mate curvature. In Francis R. Bach and David M. Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, volume 37 JMLR Workshop and Conference Proceedings, pp. 2408-2417, JMLR.org, 2015.
Tom Minka, Bayesian linear regression. July 2000.
13

---

Published conference paper at ICLR 2023
George Nicholson, Marta Blangiardo, Mark Briers, Peter Diggle, Tor Erlend Fjelde, Hong Ge, Robert J. B. Goudie, Radka Jersakova, Ruairidh E. King, Brieuc C. L. Lehmann, Ann-Marie Mallon, Tullia Padellini, Yee Whye Teh, Chris Holmes, and Sylvia Richardson. Interoperability of Statistical Models Pandemic Preparedness: Principles and Reality. Statistical Science, 37(2): 183 206, 2022. doi: 10.1214/22-8 STS854.
Søren Feodor Nielsen. The stochastic EM algorithm: estimation and asymptotic results. Bernoulli, 6 (3):457 - 489, 2000. doi: bj/10816167011
Roman Novak, Lechao Xiao, Jiri Hron, Jaehoon Lee, Alexander A. Alemi, Jascha Sahl-Dickttein and Samuel S. Schoenholz. Neural tangents: Fast and easy infinite neural networks in python. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia. April 26-30, 2020. OpenReview.net, 2020.
Roman Novak, Jascha Sohl-Dickstein, and Samuel Schoenholz. Fast finite width neural tangent kernel. In International Conference on Machine Learning, pp- 17018-17044, PMLR, 2022.
Ian Osband, John Aslanides, and Albin Cassirer. Randomized prior functions for deep reinforcement learning. Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa- Bianchi, and Roman Garnett (eds.). Advances in Neural Information Processing Systems 31: Annual Conference Neural Information Processing Systems 2018, NewrIPS 2018, December 3-8, 2018, Montréal, Canada, pp. 8626-8638, 2018.
Tan Osband, Zheng Wen, Mohammad Asghari, Morteza Ibrahimi, Xiyuan Lu, and Benjamin Van Roy. Epistemic neural networks. CoRR, 0obs/2007.08922 2021.
Tan Osband, Zheng Wen, Seyed Mohammad Asghari, Vikranth Dwaracherla, Xiuyuan Lu, and Benjamin Van Roy. Evaluating high-order predictive distributions in deep learning, In The 38th Conference on Uncertainty in Artificial Intelligence, 2022.
Yaniv Ovadia, Emily Fertig, Balaji Lakshminarayanan. Sebastian Nowozin, D Sculley, Joshua Dillon, Jie Ren, Zachary Nado, and Jasper Snoek. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. Advances in Neural Information Processing Systems, pp. 13969-13980, 2019.
George Papandreou and Alan L. Yuille. Gaussian sampling by local perturbations. In John D. Lafferty, Christopher K. I. Williams, John Shawe-Tayle Richard S. Zemel, and Aron Culotta (eds.), Advances in Neural Information Processing Systems 23: 24th Annual Conference on Neural Information Processing Systems 2010. Proceedings of a meeting held 6-9 December 2010, Vancouver, British Columbia, Canada, pp. 1858-1866, 1866. Curran Associates, Inc., 2010.
Tim Pearce. Felix Leibfried, and Alexandra Brintrup. Uncertainty in neural networks: Approximately bayesian ensembling. In Silvia Chiappa and Roberto Calandra (eds.), The 23rd International Conference on Artificial Intelligence and Statistics, AISTATS 2020, 26-28 August 2020, Online [Palermo, Sicily, Italy]. volume 108 of Proceedings of Machine Learning Research, pp. 234-244. PMLR, 2020.
Carl Edward Rasmussen and Christopher K. I. Williams. Gaussian processes for machine learning. Adaptive computation and machine learning. MIT Press, 2006. ISBN 022211853X..
Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable laplace approximation for neural networks. In 6th International Conference Leariing Representations, ICIR 2018, Vancouver BC, Canada, April 30 May 2018, Conference Track Proceedings. OpenReview.net, 2018.
Daniel E. Runcie, Jiayi Qu, Hao Cheng, and Lorin Crawford. Megalmm: Mega-scale linear mixed models for genomic predictions with thousands of traits. Cenome Biology, 22(1):213, Jul 2021. ISSN 1474-760X. doi: 10. 10.1186/s13050 02-599021-00 -244166..
Michael E. Tipping, Sparse bayesian learning and the relevance vector machine. J. Mach, Learn. Res., 1:21 1-244, 2001.

---

Published conference paper at ICLR 2023
Michael Tipping and Anita C. Faul. Fast marginal likelihood maximisation for sparse bayesian models. In Christopher M. Bishop and Brendan Frey (eds.). Proceedings of the Ninth Interna- fional Workshop on Artificial Intelligence and Statistics, AISTATS 2003, Key West, Florida, USA, January 3-6, 2003. Society for Artificial Intelligence and Statistics, 2003.
Malte Tölle,
Laves, and Alexander Schlaefer. A mean-field variational inference approach to deep image prior inverse problems in medical imaging. Mattias Heinrich, Qi Dou, Marleen de Bruijne, Jan Lellmann, Alexander Schlaefer, and Floris Ernst (eds.), Medical Imaging with Deep Learning, 7-9 July 2021, Liibeck, Cermany, volume 143 of Proceedings of Machine Learning Research, pp. 760-760. PMLR, 2021.
Dmitry Ulyanov, Andrea Vedaldi, and Victor S. Lempitsky. Deep image prior. Int. J. Comput. Vis., 1867-8888 2020. doi: 1007/| 1263-020-0130--
James T. Wilson, Viacheslav Borovitskiy, Alexander Terenin, Peter Mostowsky, and Mare Peter Deisenroth. Efficiently sampling functions from gaussian process posteriors. Proceedings of the 37th International Conference on Machine Learning. ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 10292-10302. PMLR, 2020.
David P. Wipf and Srikantan S. Nagarajan. A new view of automatic relevance determination. In John C. Platt, Daphne Koller, Yoram Singer, and Sam T. Roweis (eds.), Advances in Neural Information Processing Systems 20, Proceedings of the Twenty-First Annual Conference on Neural Information Processing Systems, Vancouver, British Columbia, Canada, December 3-6, 2007, pp- 1625-1632. Curran Associates, Inc., 2007.
Arnold Zellner. On Assessing Prior Distributions and Bayesian Regression Analysis with Prior Distributions, volume of Studies in Bayesian Econometrics and Statistics, pp. 233-243. Elsevier Science, 1986.

---

Published a conference paper at ICLR 2023
A
RELATED WORK
Bayesian Gaussian linear models This work builds on the rich literature of Bayesian linear regression (Gull, 1989; Bishop, 2006; Rasmussen & Williams, 2006). Specifically, we present stochastic approximation to the iterative algorithm for hyperparameter selection introduced by (Mackay, 1992) and extended by Tipping (2001); Tipping Faul (2003); Wipf & Nagarajan (2007); Antorán et al. (2022c). Analytical tractability makes linear models ubiquitous in machine learning, with applications in genomics (Runcie al., 2021), reinforcement learning (Ash et al., 2022), and pandemic modelling (Nicholson et al.. 2022). among others. Alas, linear models are held back by a cost of inference cubic in the number of parameters when expressed in primal form, or cubic in the number of observations for the dual (i.c. kernelised or Gaussian Process) form. Additionally, for non- non-Gaussian likelihoods, e.g. in classification, inference is no longer closed form. The most common approximations used in these settings are Laplace's 's method (Mackay, 1992) and variational inference (Hensman et al., 2013). Khan et al. (2019) and Adam et al. (2021) show that every Gaussian approximation corresponds to the true posterior of a surrogate regression problem with the same features, a fact which we use in this work to apply sample-then-opimiss to Laplace posteriors.
Sample-then-on theoopise Papandreou & Yuille (2010); de G. Matthews et al. (2017) phrase sampling from a conjugate Gaussian-linear model as solving perturbed quadratic optimisation problem. This method has been applied for uncertainty estimation in non-linearised NNs by Osband et al. (2018; 2021), and Pearce al. (2020). although in these setting it does not draw exact posterior samples. In this work, we show sample-then-optinise to be the primal form of Matheron's rule (Journll & Huijbregts, 1978; Hoffman & Ribak, 1991), method for updating jointly Gaussian samples into conditional samples, which was recently repopularised by Wilson et al. (2020).
Lincarised neural networks Introduced by Mackay ( (1992), this approximation yields closed- form errorbars for Laplace posteriors. Lawrence (2000) and Ritter et al. (2018) found the Laplace approximation to underperform without the linearisation step. Khan et al. (2019) and Immer et al. (2021b) re-popularised the linearisation step by showing that it improves the quality of uncertainty estimates. Kristiadi et al. (2020) show that the Laplace approximation is sufficient to resolve certain pathologies of soint-estimate NNs' predictions. Immer et al. (2021a) and Antorán et al. (2022a:c) explore the linear model's s evidence for model selection. Immer et al. (2022) shows the objective can even be used learn invariances in deep models. Daxberger et al. (2021b) and Maddox al. (2021) introduce subnetwork and finite differences approaches, respectively, for faster inference with the linearised model. This line of work is also closely related to the neural tangent kernel (Jacot et al., 2018; Lee al., 2019; Novak et al., 2020) in which NNs are linearised at initialisation.
The g-prior, originally introduced by Zellner (1986), consists of a centred Gaussian with covariance matching the inverse of the Fisher information matrix. Resultantly, the g-prior ensures inferences are independent of the units of measurement of the covariates (Minka, 2000). Since then, it has extensively used in the context of model selection for generalised linear models (Liang et al., 2008; Bové Held, 201 I; Baragatti & Pommeret, 2012). In the large-scale setting, we overcome the computational intractability of the Fisher by diagonalising the g- g-prior while preserving its
property.
B MODEL EVIDENCE LOWER BOUND AND THE EFFECTIVE DIMENSION
B.1 EQUIVALENT FORMULATIONS OF EFFECTIVE DIMENSION
We begin by relating two standard forms of effective dimension, which we throughout. Starting with the form standard in the kernel-bass literature (that without an explicit d' dependence).
7 ((A MY Tr ((I - = Tr (1
(18)
d' - Tr ((A + -))
(19)
we have arrived at the form used within the -dinteso-inniinall linear modelling literature (Mackay, 1992; Wipf& & Nagarajan, 2007; Maddox al., 2020).
16

---

Published conference paper at ICLR 2023
B.2 DERIVATION OF M AS A LOWER BOUND ON THE MODEL EVIDENCE
Let pe be the Lebesgue density of N2 B. ). P = 1 and Q = =N (M + Then.
log (Y:4' log / po(Y)dP= log / po(Y) da dQ 2 log po(Y) dP
dP
dQ
(20)
dQ]
logpo(Y) ) (21)
where denotes the Ki.-divergence. Starting with the first term,
/ Po
/ n log 2nr + log detB - (Y - -000) -00)) B(Y - +00)dQ
(22)
1 / (X =00 +0)dQ,
2 I-n log 2n log detB]
(23)
2
and expanding the quadratic form.
f - -O) +0)dQ = YTBY - 2Y f000=
[ddQ - / f (TTTBOOOOQ
(24)
= =Y"BY - 2Y B + ModQ.
(25)
To handle the final integral, consider that
7 (M A')
(26)
= (M / (0- 0)(0 - 0)"dQ)
(27)
= -T (MO00) ) Tr (M / 000 aot
(28)
= -0 f 0r ModQ.
(29)
and thus
f logro(YY)cQ
1
[log detl3 12 log 2n (Y -00) - 40) 7]
(30)
=log
(31)
2
The KL between two multivariate Gaussians is standard result, yielding
D(Q[PP= log detA' + log det(M + A') - d' + 0T A'O (r + ')
(32)
log det log det(M + A²) + 0 A'O A'O 7].
(33)
2
where we used that d' Tr [A'(M + A))-1}. Putting together (31) and (33), obtain
1 1013. -M(A),
logp(Y: A') log PY() log det( (A M +1)
(34)
which is the stated result up to taking C' = log (Y))
B.3 FIRST ORDER OPTIMALITY CONDITION FOR M
Consider the derivative of M. We have,
VM(A) = 2 [V [1100 + V log det(A + M) - V log det.A]
(35)

---

Published conference paper at ICLR 2023
where we expanded log det(I M- log det(A + M) log detA. Taking the respective derivatives and setting equal to zero at A, this leads to the condition
00 (1 - (1 + M)
(36)
fost-multiplying by and applying the push through identity, we obtain
00 M + M)
(37)
For the above to hold, is necessary that the traces of both sides are equal. Thus,
|013 Tr 100 A} Tr (A M)- =
(38)
which is the stated first order optimality condition, up to cyclic permutation.
B.4 M-STEP FOR FEATURE-WISE REGULARISATION STRENGTHSS
We can leverage the primal form expression for the effective dimension given in Appendix B.1 extend the above first order optimality condition to the feature- wise regulariser setting.
Consider a sub-vector of our weight vector contiguous between the ith and jth weights written as Ozg Note that we only choose contiguous weights for notational convenience but it is not necessary do in general.
The first order condition from Appendix B.3 is satisfied for any i,j with < j we have
a |as 1 = j i- Σ [A] M) ") Vag
(39)
We assume [A]us for all S <. j. Thus, we may update the regulariser for each separate weight sub-vector as a = //://10://1
ANALYSIS OF LOSSES AND LOSS GRADIENT ESTIMATOR VARIANCES
C.I ON LOSS MINIMA
The losses and L' are strictly convex, thus confirm they have the same unique minimum, it suffices to consider the respective first order optimality conditions, VI((C) = and V//((") = 0. We have,
VI()) = o5 30400 - 8) + A(S -P),
(40)
and
VL'(() BO +A(5 A-- 80) O""((C' -)) A(Ç' -0)
(41)
(42)
Thus <= Ç' almost surely. Moreover, L'(z) = L(x) + C for all for C' constant independent of z.
To determine the distribution of note that is a linear transformation of zero-mer aero-mean Gaussian random variables, and thus itself a zero-mean Gaussian random variable. Rearranging the first order optimality condition, we find that
H BE + A0))
(43)
Thus
B(CCT]
H E BE + 100) HH
(44)
= - 18 H-1
(45)
" B -
(46)
And N' (0-0 _1 claimed.

---

Published conference paper at ICLR 2023
C.2 Loss GRADIENT VARIANCE CONDITION
Taking j Unif( (..... nJ),
the gradient estimators for the data detandependent terms of and L' are
nV]()(z)) sylli, no(a;)) B,(((1;) - Eạ)
(47)
and
j' = n V (obr))))= np(aa)) By0(xz)² -
(48)
respectively. Their variances are related
Var(g) Var(no(a;) (,(((2))===00 - ej))
(49)
Var(no(w;)) (,0(2;)2 (, + 2Cov(nó(zz) )2 no(z))"BBEJ)
Var(nó(z;))BBee, B5e;)
(50)
Var(g') Var(no(a;) Bye,) 2Cov(n0(3,) BB)la)) nø(s;) Bye,)
nofz,)
(51)
Evaluating the variance and covariance, we have
Var (nn(l,)) Bys;) nVar(("
(52)
and
(ov(nó(z,)* no(xz)² nCov("" BE),,,
(53)
and thus
Varg Varg' n [Var BE) 2Cov(1 (#"'Be"BB x BE)] =nd.
(54)
C.3
CONDITION AT CONVERGENCE
Now consider Tr A for 11 the optimum of both 1. and I.' From the first order optimality condition,
< H= ( BE A6%)
(55)
Proceeding to rearrange the condition at =.
T
Tr (E (B BE 2 "e"")"
(56)
Tr (184 (E
- BE 24 B (0 BE A0"))²
(E
(57)
'r BE T BE 24 B + AE AE)
(58)
Tr (B4 (E BB BE 24 10
(59)
Tr (4" BEEE (B3 2
(60)
Tr Ba 2H B4))
(61)
where substituted in the definition of then used that & and 6 independent, and that E|g0] = 0, and finally that E[Eε') " B
Writing M BB and recalling that // (M + A), we have
T (M - 2(M + M))
(62)
[T - 2(4-' + 2
(63)
= =Tr(M(I - 2(1 -AA- + +))
(64)
-Tr(M(-I 2(M+ A)-
(65)
= (M) 2Tr A) A).
+
(66)
where we have used that (A-'M + A
- = I (A-'M 1) for the fourth equality.
Now consider the isotropic prior case A = al and recall the effective dimension is written as 7 Tr (M A)-1). The above implies Tr A > 0 if and only if 2a7 Tr Bd.

---

Published conference paper at ICLR 2023
C.4 ANALYSING CONDITION AT CONVERGENCE
To gain some intuition for the condition at convergence, denote by , the eigenvalues M (with multiplicity). We can use these to restate the condition as
207
2a
N = Tr(M)
(67)
1=1
This formulation of effective dimension gives an interpretation dimensions for which Xj is larger than a; in that sense, A measures how well determined the corresponding dimension of the weight vector / is by the observed data. From here, note that
soft count of the number
200y
min(2y.
(68)
Ag +
and thus it is sufficient for Tr A > 0 to hold at convergence that a > X, for all j (but, course, not necessary), yielding the intuition that L' is preferred when the problem is heavily regularised.
D DUAL FORM OF THE
LOSS: MATHERON'S S RULE
Both losses L and L/ result in a random variable Ç ~ IT given by
< A6%).
(69)
Recalling that H = A A+ JT ( B& and using the push-through identity, we can express equivalently as
< H ((H x BBB BEBE)
(70)
8 + H- -88)
(71)
1 ( - B(E +00%)
(72)
0 -
400
(73)
0 4- 04
+00)
(74)
Now taking sample of the posterior Gaussian process evaluated imput a to be = o(x)( and the corresponding sample of the prior process to be Go <xx)0" premultpllyin the above expression by ((z) we obtain
G= + +(z))=
$A-14" - -
(75)
which is Matheron's 's rule.
RESOLVING FEATURE SCALE INDETERMINACIES IN THE NN JACOBIAN WITH THE G-PRIOR
Antorán al. (2022c) show that for NNs with normalisation layers, the Jacobian features d(-) V.. (w, -) corresponding to each NN layer are scaled arbitrarily. To illustrate this, we divide the NN linearisation point into the concatenation of two weight vectors v [wp, W1]. We assume the layer containing Wp followed by normalisation layer, but not that containing W1, which leads to the invariance
[(kwo. wil.-) [[Wo, w1).
(76)
for all >
While f is invariant to this scaling, the Jacobian feature embeddings o(-) = v :) are not. We separate the embeddings as
((x) [do(.). (-), pi (-)] IV m. ( Vu, S(w,
(77)
Antorán et al. (2022c) show that, given a reference pair ([wo, Wi], [du(z), oi (z))), and for Wo normalised, scaling wp by & results in the pair (Rkwg, Tz], [ dn (x), or (x)]). )]. Thus. using a single prior precision parameter, the regularisation strength applied the weights multiplying Φo(x)) relative
20

---

Published conference paper at ICLR 2023
to those multiplying @1 (z) will increase with The value of k. the scale of the linearisation point. depends on exogenous factors such as learning rate or batch size-and importantly is independent of the data, since it does not affect the output.
One way resolve this is to assign the weights Wo and 11 separate regularisation parameters and learn these using the EM procedure outlined in Section 2. However, instead, consider using the g g-prior normalised features ø' introduced in Section 4.3, and specifically, the scaling vector corresponding normalised and non-normalised components and for Lh normalised, the k-scaled pair is (Ikwo, Wh and
[so. $1]- For reference pair ([wo, W1]. (So. $1))
[diag(k 05 B0 11000,10000 diag( ( (00) -0 [Kiso 81]
where @ denotes an elementwise power: Since the k-scaled features are [A: 1b-1 (+) (-)). when applying the g-prior normalisation we recover feature vector independent of k. This resolves the aforementioned pathology.
F
A PRACTICAL IMPLEMENTATION OF SAMPLE-BASED INFERENCE AND HYPERPARAMETER LEARNING FOR LINEARISED NEURAL NETWORKS
Algorithm| provides a high level overview of the procedure used for our experiments. This appendix expands on this, providing fully detailed algorithms for both sampled linearised Laplace applied to classification networks and the kernelised version of the method that use for tomographic image reconstruction.
Image classification Algorithm 2 provides an algorithm for lincarised Laplace inference using the stochastic E.M iteration presented in Section 3 for hyperparameter selection and the g-prior normalisation described in Section 4.3. Therein, / denotes the softmax function. The curvature of the cross entropy loss at zú, denoted B, is given by B, = diag(p.) - p&p. for p; f xx)) our neural network's ss predictive probabilities. The notation @ refers to the elementwise product and to the elementwise power when used in an exponent. We refer to the Cholesky factorisation of positive definite matrix as its 1/2th power.
In order to limit computational cost, we sample the stochastic regularisation terms (07) ). per (7), only once at the start. Not resampling these each E step results in the optima of the sampling objective being close for successive iterations. This comes at the cost of small bias in our estimator which we find to be negligible in practise. We separate (0") into a sum consisting of a prior sample from (0?) and data dependent term, denoted (0(). The former scales with - while the latter with a so this allows to update each term in closed form each time a changes. We initialise our samples (09) at the first EM iteration. We warm-start the posterior mode at the previous solution between iterations, initialising it to zero for the first iteration. We estimate the g-prior scaling vector s by noting that it relates to 0g as a= (1800) @ 011)0-1/-
We optimise both our samples and posterior mean Ø using stochastic gradient descent with Nesterov momentum parameter of 0.9. We find that Polyak averaging is very effective at reducing gradient variance when optimising the sampling objective (per Dieuleveut et al., 2017). However, it has two limitations 1) it slows down optimisation, increasing the number of steps needed 2) it doubles the memory requirement needed to store posterior samples. This decreases the number of samples that can be optimised in parallel on a single hardware accelerator: Instead we employ a linear learning rate decay schedule, which we find to work nearly as well while not increasing computational burden. The regularised classification loss is non-quadratic and thus Polyak averaging no longer optimal (Bach, 2014). Thus here we also employ a linear learning rate decay schedule. For optimising both the sampling and classification loss objectives we find that gradient clipping helps prevent oscillations at the start of training and as a result speeds up convergence.
The key perraarametrrs of our algorithm are the number of samples to draw for the EM iteration, the number of EM steps run, and SGD hyperparameters: learning rate, linear decay rate, number of steps, batch-size and gradient clipping. Empirically, we find that most EM steps are necessary for hyperparameter convergence and that as little as 3 samples can be used for the algorithm without degrading performance. Choosing SGD hyperparameters is more complicated. However, we are aided by the fact that lower loss values correspond to more precise posterior mean and sample estimates.
21

---

Published conference paper at ICLR 2023
Algorithm 2: Sampling- based linearised Laplace inference for image classification
Inputs: Linearised network h, linearisation point w, observations #i.... En. negative
log- log-likelhhood function initial precision a > 0, number of samples
Function B(i):
P: - (1((110 x.)) return diag(pt) P.p
for
k do I....,
09
(10,a'1))
8 +
Ei d(x)) sy N(O, -(
G 5-
0g
ト 0002 3-1/2
while a has not converged do
for j = l,...., k do
L G SGD: (1 + 09 @ (,00)) ini 0 <- SGD, EE ((x00))
+a init=0)
+ $ Ej=1 E (I(s) B()
a' - 1/1001
for = I,....k k do
8 t 0
a' Output: Optimised precision a and weight samples ÇI...., Ck
Algorithm Kernelised sampling-basee based linearised NN inference for CT reconstruction Inputs: Linearised network h, linearisation point w, measurements Y, discrete Radon transform U, U-Net Jacobian 4, initial precision 0, number of samples l noise precision B
Function Kvp U B- return (
Eien (2) 1/2
while a has not converged do
P + Compute-precondiioner(sss for = 1...., & do
Ç < Ud(s @ 89) + E where E, ~N(,00 and 09~~ ~N N-00A-1))
Cj < CG (Kvp, G., precond.=PP
Çj + <3 - (v(a diag( s (20))=TTT
8 + U(DW f(w))
C - CG (Kvp, Y+ precond.=P)
0 - 8 - TT
3 ト 1 و
(3)
0/01/10012 a+
Output: Optimised precision
22

---

Published conference paper at ICLR 2023
As a result, we can tune these parameters on the train data, no validation required. The specific hyperparameter values used in our experiments are provided in Appendix G.
A final thing to note is that due to the presence of normalisation layers and a dense final layer. for classification networks, the constant-in-6 terms cancel in the linearised model and are left with ar) = <(x)0 (Antorán et al., 2022c). In our algorithm, this fact is only relevant to the computation of the posterior mode 0 as the optima of ())(
Tomographic reconstruction Algorithm 3 is the kernelised version of algorithm that we for tomographic reconstruction, This problem described in detail in Appendix G.3.
Distinctly from the image classification setting. tomographic reconstruction is a regression problem for which we a Gaussian likelihood with fixed noise precision B J. The linear model's loss function L is thus quadratic and the Laplace approximation is not needed. Both the sample loss and the linear model's 's loss can be optimised in closed form by solving a linear system of equations given by the observation covariance, i.e. the kernel matrix, (1 - "diag( (s ( o0t U + B where the linear operator U represents the discrete Radon transform and combines with the U-Net Jacobian build the feature embedding U/
Wc solve against the kernel matrix using the preconditioned conjugate gradient (CG) method de- scribed by Gardner et al. (2018). As preconditioner, we compute 400-dimensional randomised eigendecomposition (alg. 5.6 in Halko et al., 2011) preconditioner, which we invert using the Woodbury identity. We find the preconditioner to provide important speedups to CG convergence and we re-estimate it after every hyperparameter update. Both computing the preconditioner and running preconditioned CG optimisation only interact with the kernel matrix by computing its products with vectors. Our algorithm defines our kernel vector product Kvp routine explicitly, as it is central to our implementation. We find that the GPyTorch CG implementation does not benefit from warm-starting the solution vector. Consequently, we re-draw prior and noise samples (8") E) every E-step.
Similarly to image classification, the key hyperparameters are the number of samples to draw for the EM iteration, the number of EM steps to run, and CG optimisation hyperparameters. Again, the number of samples can be kept low (e.g. 2) and we find around 5 steps to suffice for convergence of the prior precision a. The key conjugate gradients hyperparameters are the tolerance at which to stop optimisation and the maximum number of optimisation steps if the tolerance is not reached. We provide our choices in Appendix G.3 but note that our of large preconditioner results in CG always hitting the desired low error tolerance within steps and never stopping due to reaching the maximum number of steps. In turn, this makes our kernelised EM algorithm notably faster than its primal form SGD-based counterpart.
particularity of this setting is that the U-Net does not have a dense final layer. As a result, the constant- in-0 terms in the linearised function do not cancel (see Section 4.1), leading to the appearance of the target offset term 8 when solving for the posterior mean.
G EXPERIMENTAL DETAILS
In this appendix we provide experimental details and hyperparameter settings omitted from the main text.
G.1 MNIST EXPERIMENTS
MNIST m=10 way classification experiments were performed using the LeNet-style CNN archi- tectures of increasing size employed by Antorán et al. (2022c): "LeNetSmal"" (d'=14634), "LeNet" (d'-29226) and "LeNetBig" 460=4.0. We note that these models contain batch normalisation layers. Each model was trained with using SGD with momentum of 0.9 for 90 epochs with learning rate drop of a factor of 10 every 30 epochs. The MNIST dataset was downloaded from PyTorch torchvi sin. We employ standard per-channel mean and std-dev standardisation preprocessing and two pixel shift and crop data augmentation. For posterior mode optimisation and sampling, we do not perform data augmentation as to avoid cold posterior effects (Izmailov et al., 2021). The details of SGD approaches to convex optimisation for obtaining posterior modes and samples as follows
23

---

Published conference paper at ICLR 2023
Posterior mode optimisation: The linearised NN weights are trained using SGD with a Nesterov momentum coefficient of 0.9, and batch size 1000 for 40 epochs. We clip gradients to maximum norm of 1. We use an initial learning rate of Ie - 2 when using standard isotropic or layerwise Gaussian priors, and for the g-prior. We employ a linear decay schedule that reduces the Ir by a factor of 330 over the first 75% of the training procedure and holds it constant afterwards.
Sampling: optimise 32 samples in parallel using SGD with Nesterov momentum (=0.9) and a batch size of 1000 for 20 epochs, For standard Gaussian priors (isotropic and layerwise), we learning rate of 1 whereas for the g-prior, we find higher learning rate of 200 to work best.
Hyperparameter optimisation: We tuned the learning rate, decay schedule and gradient clipping strength using rough grid search over multiple orders of magnitude. We chose the settings that reached the lowest loss values. These can be evaluated with just the train set. We chose the largest batch size that could accommodate optimising 32 samples in parallel on a single hardware accelerator. We note that posterior mode and sample optimisation converge in less than half of the total epochs we use for their optimisation. The numbers of epochs chosen were set to be large enough to ensure convergence and not tuned. A decrease in computational cost can likely be achieved by stopping sample optimisation earlier.
Baseline methods. For the comparison learning single precision hyperparameter and layerwise hyperparameters in Figure 3, we extend the M-step update to as oq = 1//16:11 where / indexes each layer's 's attributes, as done in (Mackay, 1992; Tipping, 2001). For the MAP, diagonal covariance and KFAC covariance baselines, we use the same pre-trained models when possible (i.e. not for the ensembles or dropout baselines). Since all baselines share the same linearisation point, they also share the same mean predictions. Differences in performance among baselines are thus only due to differences in uncertainty estimation. The diagonal approximation the covariance is constructed by first computing the diagonal of the Hessian M and the inverting it. For the KFAC covariance approximation, exploit the equivaleney between the Generalised Gauss Newton matrix (i.c. the Hessian of the linear model h) and the Fisher information matrix for exponential family likelihoods (i.e. the categorical). This allows us to formulate the Hessian as an expectation of likelihood gradients, which in turn we approximate using single sample per training observation, as in (Daxberger al., 2021a). For completeness, we also state the probit approximation for sampled predictive posteriors over logits. For input æ and samples Ci...,, Se. the predictive probability for class i E [|l..... m|| is given by
[(W,)) (060))))22)0003
softmax
2h
G.2 CIFAR 100 CLASSIFICATION
CIFAR100 m=100 100 way classification experiments were performed using ResNet18 mod- els (d' z 11M) with specific architecture details matching the PyTorch torchvision implementats iinn. We train these models using SGD with momentum of 0.9 for 300 epochs. The starting Ir is 0.1 and we reduce it by factor of 10 every 100 epochs. The CIFAR 100 dataset was also downloaded using torchvis sion and our data preprocessing and augmentation also follow the default implementation from this library. For posterior mode optimisation and sampling, we do not perform data augmentation. The SGD details used to solve the convex optimisation problems required for obtaining posterior modes and drawing samples are as follows
Posterior mode optimisation: The linearised NN weights are trained using SGD with Nesterov momentum (=0.9) and batch size of 2000 for 40 epochs. We employ linear decay learning rate schedule with an initial learning rate of le-1. It is decreased by factor of 330 over the first 75% of training. and then held constant. We also employ gradient clipping with maximum norm= 0.
Sampling: We optimise samples in parallel using SGD with Nesterov momentum (=0.9) and batch size 100 for 20 epochs. All other details match those of posterior mode optimisation.
24

---

Published conference paper at ICLR 2023
Upon convergence the EM algorithm, we draw 64 further samples using the optimal prior precision by following the optimisation procedure described above. We initialise these samples at prior samples drawn with the optimised prior precision.
Hyperparameter optimisation: We tuned the learning rate, decay rate and gradient clipping strength using a rough grid search over orders of magnitude. We also chose the largest hatch size that for which we could simultancously optimise 6 samples in parallel on single hardware accelerator. Similarly to the MNIST experiments, instead chose large values that would ensure convergence. sped up by decreasing the duration of the convex optimisation routines.
did not optimise the number of optimisation epochs and
likely that our EM iteration can be
Details for baselines and hyperparameters not mentioned explicitly in this subsection match those given for MNIST in the previous subsection.
G.2.1 EFFICIENT K-ADIC SAMPLING
Osband et al. (2022; 2021) introduced dyadic test input sampling (k 2) as practical way of evaluating joint predictions discriminative tasks. This method samples K random anchor points from the test dataset, and then randomly resamples them to create a batch of size 7 = 10. Test log-likelihood is evaluated jointly for each batch as
log
exp
Σ elv, f(0,x,0 dII, iST
for f the model being evaluated and II its posterior distribution over model parameters. This quantity be estimated with posterior samples Se as
log
cxp
Σ
Eeenn f( F5
JSk
We extend this evaluation approach to larger i and values without increasing computational cost. We randomly sample integers (bi,...,b.) such that they sum to i.e Σ k The joint log-likelihood over the batch size 7 with x unique datapoints can then be estimated as
log Σ exp
Σ bel(on, f(C), zi))
JSk
where the inner sum is over the distinct elements the batch instead of the "total batch size" This is equivalent to the formulation proposed in Osband et al. (2022) for dyadic sampling, when x = 2 and 7 = 10. We note that it is not possible to achieve augmented dyadic sampling, as described in (Osband al., 2021), with this approach. However the authors mention that there is not a significant difference in the relative performance of methods when using augmented dyadic sampling compared regular dyadic sampling. We introduce a final step however, which is to repeat the computation for multiple shuffles (10) of the test dataset. This eliminates variance our results from the choice of the i observations which get grouped together in each batch.
G.3 TOMOGRAPHIC RECONSTRUCTION
Problem setup Tomographic reconstruction consists in solving linear inverse problem in imaging
where we observe a set of measurements E R'" which we assume to be generated as y y= Uxx y=Ux"* +
for U e R the discrete Radon transform, x² E Rd the image to reconstruct and 1) ~ (000, I) random noise. We have < d, making the problem underconstrained. Our specific tomographic reconstruction task closely follows the one from Barbano et al. (2021). We perform a sperse-view reconstruction of an image of a slice of a walnut from a sub-sampled set of measurements. Specifically, from the full measurement equidistant angles over [0, 360%). ). we choose our measurement set by subsampling angles by a factor of either 10x or 20x, leading to measurements of size m 15360 or m 7680. As in Barbano et al. (2021); Antorán et al. (2022b); Barbano al. (2023), we reduce the original 3D scan geometry the 2D slice interest by selecting the relevant subset of measurement pixels. We assemble the Radon operator U as a sparse matrix taking in an image of resolution (501px)² and outputting a measurement tensor coherent with the descrihed geometry.
of (Der Sarkissian et al., 2019a), which containing scans 1200
25

---

Published conference paper at ICLR 2023
Image reconstruction predicted empirical
CIFAR100 predicted probability vs empirical accuracy
satreled In. UNNtt
Iin UoNet
U-M
08
emptrical coverage
0.4
0.7
MAP
Ensumble
Sampling MC
00
0.0
0.2
03.
0.6
0.8
30
0.0
02
0.6
1.0
posterior credible interval
predicted probability
Figure 8: Left: empirical coverage of test targets for posterior credible intervals of increasing width for our U-net tomographic reconstruction experiment (Section 5.3). Right: confidence vs accuracy plot (also known reliability diagram) for our CIFARI00 classification experiment (Section 5.2).
Methods To provide a reconstruction, we use the Deep Image prior (Ulyanov et al.. 2020) which trains the parameters w Ied of a fully convolutional U-Net autoencoder f IRd" Rd, where the input is fixed and thus absent from our notation, until a satisfactory reconstruction ((ū) is obtained. The U-Net network architecture is the one proposed in (Baguer et al., 2020). The optimisation of the U-Net parameters follows Barbano al. (2021), although we note that faster optimisation strategies exist (Barbano et al., 2023). To estimate the uncertainty in this reconstruction, we linearise the U-Net around v, as described in Section 1. This leaves with model affine the parameters and with design matrix U R=xdd" We may now proceed with lincar model inference. While (Antorán et al., 2022b) use the traditional EM iteration described in Section 2, we use the sample-hased one from Section 3. For evaluation, we use the non-sparse reconstruction (using 1200 angles) provided by (Der Sarkissian al., 2019a) the ground truth image x". To evaluate joint log-like le-likeli we estimate the predictive covariance matrix for patches of neighbouring pixels using samples. Covariance matrix estimates from samples are known to be unreliable. We use the stabilised formulation of (Maddox et al., 2019): Σ
24 - 23 + 25 for (2z) samples from the predictive posterior over patch. We then construct predictive distributions over pixels as N(f(w), E).
Hyperparameters
We employ a low CG tolerance of 1e 3 and a maximum number of iterations of 150, which is never reached in practise the error tolerance level is always hit in less steps.
H CALIBRATION OF PREDICTIVE DISTRIBUTIONS
This appendix evaluates the calibration of the predictive distributions provided by the methods under consideration our CIFAR100 classification experiment (Section 5.2) and U-net image reconstruction experiment (Section 5.3).
For classification, separate our predicted probabilities into 10 equal width bins between 0 and 1. For each bin, we plot the proportion of targets that coincide with the class for which the predicted probability falls into the bin. This is shown on the right hand side of Figure 8. Consistent with the results from the main text, KFAC' overestimates uncertainty all confidence levels whereas MAP underestimates it. Both sample-baset linearised Laplace and ensembling show significantly improved calibration. While ensembles show a small amount of uncertainty overestimation consistently, our method underestimates uncertainty for low predicted probabilities and overestimates it for large predicted probabilities.
For image reconstruction regression, we first compute normalised residuals by subtracting our predictions from the targets and dividing by the predictive standard deviation. Our predictive distribution for these normalised residuals the centered unit variance Gaussian. We consider
26

---

Published conference paper at ICLR 2023
posterior credible intervals centered and of increasing width plot the proportion of test points that fall within them in the left side plot of Figure 8. We find dropout inference to underestimate the magnitude of the residuals across all credible interval widths. Linearised inference with a single EM step, (Antorán al., 2022b), consistently overestimates uncertainty. Our approach, which performs 5 steps of EM, overestimates uncertainty, but to a much smaller degree, presenting the best overall calibration. The latter two approaches consist of the same model but with different regularisation strength. The difference between the two reveals the paramount importance of tuning the prior precision hyperparameter well.
I ADDITIONAL EXPERIMENTS
This appendix contains additional experiments and baselines that supplement the experimental results provided in the main text.
L.1 COMPARING PRIMAL DUAL EFFECTIVE DIMENSIONALITY ESTIMATORS
Our main-text experiments employ the kernelised effective dimension estimator introduced in (5). A different unbiased estimator may be obtained in primal form following the derivation provided in Appendix B.1. Figure 9 compares both estimators when applied to the 1d toy problem used generate Figure from the main text. particular, we use hidden layer MLP with layemoor after every hidden layer and the "Matern" dataset of Antorán et al. (2020). We use 8 samples from the exact linearised Laplace posterior to compute effective dimension estimates and repeat this procedure 1000 times to characterise the behaviour of each estimator. effective dimension using eigendecomposition.
reference, we also compute the exact
Both estimators present distributions centered at the true effective dimension value. However, the prediction space (kemellsedd estimator presents much lower variance of 9. opposed 654.19 from the weight space estimator. Additionally, the weight space estimator distribution places a substatial amount of probability mass negative effective dimension values. From the form of (19), that this is due to 8-sample estimator overestimating posterior variance. On the other hand, the kernelised estimator in (5) can only produce positive values.
eff dim estimator with K-8 samples
10-1
density
10022
10-3
50
25
125 150
eff dim estimate
exact eff dim
weight space estimakur var: 654.19
pred space estimator
var: 9.16
Figure Histogram, bin heights normalised to represent density estimates, of effective dimension estimates produced by the primal form (weight space) estimator (19) and the kernelised (prediction space) estimator (5). Both distributions roughly centered the true effective dimension but the kernelised estimator presents much lower variance.

---

Published conference paper at ICLR 2023
1.2 EVALUATING APPROACHES TO UPDATING EYPERPARAMETERS IN THE M-STEP
This section empirically motivates the exed-point iteration M-step introduced by Mackay (1992), and described in Section 3, by comparing it with alternative approaches to updating hyperparameters. In particular. we compare Mackay's 's update with the standard Laplace M-step evidence, denoted M and given in (3) and Appendix B.2, and a Gaussian ELBO with optimised mean and covariance. The latter two objectives differ in that the regulariser appears inside of the log- pog-determinant term in M. while the ELBO's covariance does not change with the regulariser while performing the M-step. Both of these objectives differ from the Mackay update that they provide an objective which requires gradient-based optimisation in the M-step. Instead, the Mackay update has closed-form.
Mistep objective comparison
M-step convergence comparison
Exact eviience
1000
HURO
Mackey updale
16P
α
steps
Figure 10: Left: exact linear model evidence for a linearised hidden layer MLP with layer normalisation together with the lower bound presented in (3), M,, and ELBO where the Gaussian posterior covariance is decoupled from the regulariser. All curves use an initial regulariser a 500 and have marker placed at their optima. Right: values of the regularisation strength obtained at successive EM iterations while using the different update strategies under consideration for the M step. Note that when we assume access to the exact evidence function, the regulariser converges a single step and no EM iteration is necessary.
Mackay M-step
M M-step
ELBO M-step
230
-300
400
430
109
10a
10P
10
106
10P
161
102
101
10a
101
102
108
104
C
evidence
step2
step
slep
step
Exact linear model evidence for a linearised 2 hidden layer MLP with layer normalisation together with the lower bound presented in (3), M (left and middle plots), and an ELBO, where the Gaussian posterior covariance is decoupled from the regulariser (right hand side plot), at different EM steps. We update the regularisation strength with Mackay's 's fixed point iteration for the left side plot. Note that M curves are shown in this plot. We maximise M in the middle plot and maximise the ELBO in the right hand side plot. All curves use an initial regulariser of a = 5 and we place a vertical dashed line at each step's update.
Figure 11:
The plot the left of Figure compares the exact lincarised Laplace evidence for hidden layer MLP with layernorm trained on the toy dataset of Antorán et al. (2020) with the bound M (3) and the decoupled ELBO. The initial regulariser is set to = 500. The ELBO is only tight for regulariser values very close initialisation, resulting in very small M steps. M is tangent the evidence the same point as the ELBO but presents a much better approximation as we move away from & = 500, The optimum of M is much closer to the optimum of the evidence. The Mackay update does not lower bound but instead provides an updated value for a which is even closer the optimum of the evidence. The right hand side plot shows the change in the regularisation parameter across successive M- M-steps using the update methods under consideration. The Mackay M-step converges the optima of the evidence in 2 steps. Using M as objective results in convergence after 5
28

---

Published conference paper at ICLR 2023
steps. On the other hand, the ELBO update requires around 100 steps. Figure further illustrates hyperparameter learning in the 1d toy setting by showing the successive lower bounds obtained by each of the approaches under consideration at each M-step. Interestingly, the Mackay update produces regulariser updates that almost exactly maximise M.
Figure compares the evidence lower bounds of the form M, given in (3), when using different covariance matrix approximations in the MNIST classification setup presented in Section 5.1. In particular, consider the full-covariance Laplace evidence. which note does match the exact model evidence due to the non- non-quadrenn quadratic classification loss, the KFAC approximation to the covariance (labelled KFAC GGN), single-sample KFAC Fisher estimate of the covariance, the KFAC empirical Fisher matrix (Immmr et al., 2021a), and diagonal Laplace covariance. We also include sample estimate of the ELBO described above. all cases, we initialise the regulariser at an optima found by applying the EM algorithm while using the full covariance M in the M-step. In this way, may the deviation of different objectives' optima from the optima of M as estimates bias in their corresponding approximations.
Figure 12 shows the KFAC and KFAC-Fisher approximations result in a systematic overestimation of the evidence optima which grows with model size. This issue is even more pronounced for the diagonal covariance approximation. On the other hand, we find the empirical Fisher provide an accurate approximation. A similar finding is reported by (Immer et al., 2021a). This is surprising, given that the empirical Fisher is known provide heavily biased estimate of loss curvature and thus perform poorly for optimisation tasks (Kunstner al., 2019). The sample-based ELBO shows close to no bias when using 16 samples. This result agrees well with our experiments from Section 5.2, where the sample-base based EM algorithm behaves well even when using very few samples.
LeNetSmall
LeNet
LeNetBig
5000
103
M
M
diag
Fisher My
sampled ELBO samples
Figure 12: Full covariance linearised Laplace evidence M together with approximations to this curve that rely on different covariance matrix approximations. We consider convolutional networks of increasing size (left to right) trained on the MNIST dataset.
1.3 CIFARI00 CLASSIFICATION
Additional Baselines In the main text, we report the test log-likelihood obtained by our method as well as that of a point-cstimated NN (MAP), an ensemble of of point-estimate NNs (Ensemble 5), and linearised Laplace with a KFAC-approximated posterior covariance matrix (KFAC), Here, we report further comparisons with other baselines handard-in-literatuue
diagonal approximation of the Laplace covariance matrix (diag), a Laplace approximation over a selected subset the full NN weight space (subnetwork** (Daxberger et al., 2021b), and Laplace approximation over only the last-layer weights of the NN with KFAC covariance matrix approximation (KFAC-LLL*) (Eschenhagen et al., 2021). 1). Note that the last layer contains 51200 weights and thus its full Laplace covariance matrix is too large invert our A100 GPUs. We distinguish the last two methods with a star (*) denote that they require cross cross-validan validation with a held-out to tune the regularisation strength hyperparameter. In particular, we use 50 held-out points from the test set, and evaluate these methods on the remaining 9950 points. This gives these methods slight ad advantage over the approaches uncertainty quantification considered in the main text. Similarly the main
29

---

Published conference paper at ICLR 2023
text, we estimate the KFAC and sampling posterior distributions with 64 samples. We use exact marginalisation, computing full Jacobians, for the diagonal covariance approximation.
test L1
-3
KFAC
Dropout
Ensemble (5)
ding
KFAC-LII
subnetwork"
Sampling MC
Sampling-Probit
Corruption severity
Figure 13: Performance under distribution shift for additional inference baselines applied to ResNet18 the CIFAR 100 dataset. We note that KFAC-LL and subnetwork inference require held-out validation set tune hyperparameters and thus we mark them with a star (*).
Robustness distribution shift We provide the test log- likelihood results obtained with all methods under data corruption of increasing intensity in Figure 13. Following the standard setup the literature (Daxberger et al., 2021b;a; Eschenhagen et al., 2021), employ the multi-class probit approximation map Gaussian posteriors over NN outputs class probabilities for all methods except ensembles and MAP. However, when combined with our sampling approach, we find the probit approximation to overestimate uncertainty in-distribution. We illustrate this by plotting additional curve for sample-based inference with Monte-Carke marginalisation over NN outputs. This approach provides stronger in in-distribution performance which comes very close that of ensembles, subnetwork inference (*), and KFAC-LI (*). The strong performance of the latter two methods reveals the relevance of selecting good regularisation strength parameter uncertainty quantification with Laplace-style methods. In the out-of distribution setting. the probit's increased uncertainty results in larger log-likelihood scores than Monte Carlo Marginalisation. KFAC- LL performs very competitively both with ensembles in-distribution and with our approach in the out-of of-iiiiiiibttion setting.
the Gaussian distribution
Joint LL. We report marginal and joint test log-likelihood for the KFAC-LL and subnetwork inference baselines in Table 3. We use the same x-adic sampling setup in Section 5.2. marginalising the Gaussian posterior over network outputs with Monte Carlo for all methods. KFAC-LL is once again quite competitive with our approach in terms of both marginal and joint LL.
Predictive uncertainty vs number of samples. In the main text, report our method's predictive performance when drawing 64 samples. In Figure 14, we plot the degradation in test log-likelihood for the standard and progressively corrupted CIFAR100 test sets when decreasing the number samples used for prediction. We provide results for both Monte Carlo and probit marginalisation. Our results, show two main trends: 1. Monte Carlo marginalisation provides better results in-distribution for all numbers samples. This coherent with our above observation that the probit approximation results in uncertainty overestimation. On the other hand, Monte Carlo marginalisation is unbiased. 2. The probit approximation henefits less from increased numbers of samples. This is expected, since
30

---

Published conference paper at ICLR 2023
Table 3: Comparison of methods' marginal and joint prediction performance for ResNet18 on CIFAR 100. We include baselines that require validation-based tuning of the regularisation strength, marking them with star (*).
MAP
Ensemble
KFAC
Sampling
KFAC-LL
subnetwork
marginal LL
-.
0.00
-0.90 0.00
-1.12 0.01
-1.07 0.01
0.001
-1.21
2
-13.97
0.01
-6,86 0.00
-4.92 - 0.04
-5.14
-5.41 0.05
-8.38 0.07
-27.89
-14.17
-10.8
-..077
-11.15
joint L.
4 -41.88 + 0.03
-22.29 0.04
-19.02 0.22
-18.0440.18
-18.21 0.18
-25.77 0.18
5
-35.89 + 0.02
-31.07
+ 0.09
-29.40 0.40
26.75
26.50 + 0.28
-34.91 = 0.30
the probit method discards covariances in the distribution over network outputs. The probit predictive distribution has less degnées of freedom and can thus be estimated better with less samples.
Split
Split
Split 2
2.25
test LL
test LL
test L.I.
-250
-2.0
-144
-275
-2.
-3.00
48
16
36
48
64
48
4.8
16
36
48
64
Split 3
Split 4
Split 5
tost LL
test LL
test LL
-3.0
-3.5
-4.0
Sampling Probit
-3.5
Sampling MC
-45
48
16
36
48
64
4.8
16
48
64
48
16
36
48
64
samples
samples
samples
Figure 14: Predictive Performance for sample-based inference the CIFAR 100 in-distribution and corrupted while varying the number samples used for estimating the predictive distribution.
I.4 TOMOGRAPHIC RECONSTRUCTION
Stability of stochastic EM in the 15360 setting Figure 15 shows the prior precision values, effective dimension values and test log- lig-liielihood values obtained each EM step for the larger 120 angle (m = 15360) image reconstruction task. The regularisation strength converges faster in this data-rich setting than in the 60 angle setting considered in the main text (Figure convergence occurring after EM step instead of 2. We see a slightly larger sensitivity to the number of samples in this larger setting. However, the difference in test LL obtained after running stochastic EM with and 256 samples remains smaller than 0.01 nats.
with
Further analysis of uncertainty calibration The rightmost plot Figure 15 compares joint test log-likelihood obtained by our method and MC dropout on image patches of increasing size when the observation dimension is set our method performs better across patch sizes. Qualitatively, Figure shows sample-baser approach to yield uncertainty estimates with a much larger dynamic range; some pixel regions are assigned large errorbars, while others assigned small errorbars. MCDO produces less fine-graie mre-grained outputs and assigns relatively small errorbars to whole sections of the image.
15360. Similarly the results shown in the main text,
Finally, Figure 17 and Figure 18, compare the reconstruction error and uncertainty histograms for both uncertainty quantification methods under consideration for both the m 7680 and m = 15360 settings. In both plots, sample-based linearised Laplace inference slightly overestimates uncertainty.

---

Published conference paper at ICLR 2023
all. Jim
shophaal
sumplo
FM
EM
FM stegs
patcheaiz
Figure 15: Sample-based EM iteration convergence for tomographic reconstruction given = 15360. The prior precision (left), the effective dimension (middle left), and the marginal test log- likelihood (LL.) (middle right) are plotted as function of the EM step. The plot on the right shows the joint test LL across image patches of increasing size for sampling inference and an MC dropout baseline (MCDO).
1x* f(09)
sld dev
Ix" [(0)
sld dev
sampp lin UNoc
MCD0-UNNe
Figure 16: Original 501 x501 pixel walnut image and reconstruction for m- 15360 dimen- sional observation, along with pixel-wise std-dev obtained with sampling lin. Laplace and MCDO.
MCDO systematically underestimates uncertainty for the pixels where the reconstruction error largest. Interestingly, our method shows to be slightly wonsely calibrated in the more data-rich setting, the reconstruction decreases faster than the predictive standard deviation.
32

---

Published conference paper at ICLR 2023
102
IV
ded (aagp Ain UNot
107
der (NCD
density
denstyy
10
10
10-2
10-2
10-3
10-3
Figure 17: Histogram of the absolute pixelwise error computed between the reconstructed walnut image given m 7680 observations and the ground-truth for both -UNtt and MCDO-UNet. We also include the histograms of both methods' predictive standard deviations across pixels.
30
je- fill (MCDOA
(aagg
density
denaity
10-1
80-1
11-33
10
0.0
0.2
2.0
0.4
Figure 18: Histogram of the absolute pixelwise error computed between the reconstructed walnut image given m = 15360 observations and the ground-truth for both lin.-UNet and MCDO-UNet. We also include the histograms of both methods' predictive standard deviations across pixels.
33


Reviews:
[{"review_id": "m7ZD61ofFh", "reviewer": "Reviewer_d9Fu", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The computational cost associated with Bayesian linear models constrains this method’s application to small networks, small output spaces and small datasets. Moreover, the method is sensitive to the choice of regularisation strength for the linear model. This paper overcomes these limitations by introducing a scalable sample-based Bayesian inference method for conjugate Gaussian multi-output linear models, together with a matching method for hyperparameter selection. The experimental results demonstrate the contributions. ", "summary_of_the_review": "I am concern about the novelty of this paper. \n\nThe experiments are also suggested over larger models.", "strength_and_weaknesses": "Strengths:\nThe paper is well-written.\nThe approach is easy to understand and presented well in the paper. \nThe authors applied their method to some tasks to demonstrate the validity of the proposed method. \nThe theoretical analysis is strong.\n\nWeaknesses:\n1. Novelty is somehow limited. The main contribution is that the authors simply apply EM algorithm to scale inference and hyperparameter selection in Gaussian linear regression. It would be more novel if the EM algorithm was improved. \n2. Experiments across multiple datasets and with larger models would make the results more convincing. \n3. The ResNet-18 is also a very small model. Why not try the method on larger models? Is it because of the difficulty to implement?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "TI4m9Oez63", "reviewer": "Reviewer_MCAU", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper presents a number of tricks to scale linearization + Laplace approximations to large neural networks with the goal of performing uncertainty quantification.  The standard approach is to train a neural network and then post-hoc linearize the network, and use a Laplace approximation to obtain that the posterior should be Gaussian.  In this posterior, a regularization term (corresponding to the prior) appears, and this term should be fit using maximum likelihood.  To do so requires some computationally expensive numerical matrix algebra, and so the present paper instead proposes some sampling approaches (which in turn can be turned back into optimization problems) to more efficiently fit the prior and obtain samples from the posterior.  The authors apply this approach to a number of networks and datasets, finding that it can obtain better uncertainty estimates faster than existing methods (particularly in the case of simultaneously estimating uncertainty at several points).", "summary_of_the_review": "Overall, I found the paper to be a collection of interesting tricks that scale an existing approach to larger neural networks and datasets.  The presentation was mostly clear, and the presented approach seems like it will be useful.", "strength_and_weaknesses": "Strengths:\n* The approach is interesting and combines a lot of disparate ideas to make this the linearization + Laplace approximation scalable.  Some of these ideas may be of (minor) independent interest.\n* The empirical runtime and uncertainty quantification results are impressive, and show a benefit to this approach for real networks.\n* The writing was usually clear (but see below), and the results were presented clearly.\n\nWeaknesses:\n* In a few places the presentation was hard to follow.  For example, the rightmost plot in Figure 3 shows different layer-wise precisions, but the exposition only discusses learning a single precision (controlled by $\\alpha$).  I can imagine how it is straightforward to extend to different layer-wise $\\alpha$s, but that should be clarified in the main text.  In Section 4 , $\\bar{w}$ is used without definition (presumably the optimized weights).  Then, $\\bar{\\theta}$ is used (e.g., in equation (15)) without definition.  Finally, the acronym MCDO is introduced _after_ references to Figure 6 and Table 2, where that acronym is used.  Similarly \"DIP\" is used in that Figure and Table, but never defined.\n* The reported metrics are all good, but one of the most interpretable measures of uncertainty quantification is simply coverage, and it would be nice to report empirical coverage at different posterior credible levels on some held out data.  \n* I found the results in Figure 5 difficult to interpret.  Namely, it is not obvious to me how any Bayesian uncertainty quantification should perform when some datapoints are corrupted.  Would the true posterior (perhaps with a learned prior) outperform approximate posteriors in terms of log likelihood on test data?  I'm not sure that that is true, and so I'm not sure what the relative ordering of the different uncertainty quantifications methods in Figure 5 says about their \"quality\".", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Es8I17wzZ4S", "reviewer": "Reviewer_ejzP", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This work proposes a scalable linearized Laplace procedure which can be applied to realistic BNN models.  The algorithm can be viewed as an extension of the sample-then-optimize procedure that allows for normalization layers, and hyperparameter selection is also implemented a an E-step.  The authors evaluated the method on ResNet-18 and U-Net models, demonstrating improvements on predictive performance and quality of uncertainty estimates.", "summary_of_the_review": "Nice work with methodological and empirical contributions.", "strength_and_weaknesses": "The authors proposed a scalable approximation to the Laplace method, which is always of interest.  The proposed method is easy to implement, and demonstrates competitive performance.  The authors provided new connections to the sample-then-optimize formulation.\n\nAs for limitations, it is easy to nitpick (e.g., the limitation of Laplace method, the conjugate Gaussian model, or the linearized approximation), but I don't think these are important concerns that must be addressed in a single paper, or a single thread of literature, and the applicable scope of the proposed method is quite clear.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "3dKHCmnvsV", "reviewer": "Reviewer_ZeYJ", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The authors consider the task of improving the scalability of the popular linearised Laplace method for which they develop a sampling-based EM approach. The value of the proposed approach is then demonstrated in several experimental settings.\n", "summary_of_the_review": "The paper sets out to solve a well-defined task and delivers both concerning the theoretical as well as empirical side.\n", "strength_and_weaknesses": "The authors tackle an important task as the linearised Laplace is a popular approach, especially in the Bayesian Deep Learning literature. The experimental evaluation is extensive and considers both small/large scale classification tasks as well as an image reconstruction task. A minor weakness is the lack of downstream evaluation of the quality provided by the method's predictive uncertainty. However, that can be justified by the necessary limit of the breadth of the work given the page constraints.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster