S AMPLING - BASED INFERENCE FOR LARGE LINEAR MODELS , WITH APPLICATION TO LINEARISED L APLACE
Anonymous authors
Paper under double-blind review
A BSTRACT
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).
I NTRODUCTION
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 et al., 2021b; Daxberger et al., 2021a).
The method constructs a surrogate Gaussian linear model for the NN predictions, and uses the error bars of that linear model as 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 of the problem are intractable.
This restricts the method to small network or small data settings.
Moreover, the method is sensitive to the choice of regularisation strength for the linear model (Immer et 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 a samplebased 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 improve posterior sampling for large-scale conjugate Gaussian-linear models using a novel sample-then-optimize 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 us 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 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
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 CIFAR100 (100 classes × 50k datapoints) image classification using an 11M parameter
ResNet-18, and on a high-resolution (251k pixel) tomographic reconstruction (regression) task with a 2M parameter U-Net.
In tackling both of these, we encounter a pathology in the M-step of the procedure first highlighted by Antorán et al.
(2022c): the standard objective therein is illdefined 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 a standard featurenormalisation method, the g-prior (Zellner, 1986; Minka, 2000), resolves this pathology.
For the tomographic reconstruction task, the regression problem requires a dual-form formulation of our
E-step; interestingly, we show that this is equivalent to an optimisation viewpoint on Matheron’s rule (Journel & Huijbregts, 1978; Wilson et al., 2020), a connection we believe to be novel.
C ONJUGATE G AUSSIAN REGRESSION AND THE EM ALGORITHM
We study Bayesian conjugate Gaussian linear regression with multidimensional outputs, where we observe inputs x 1 , . . .
, x n ∈ R d and corresponding outputs y 1 , . . .
y n ∈ R m .
We model these as
[Equation]
where ϕ : R d (cid:55)→ R m × R d ′ is a known embedding function.
The parameters θ are assumed sampled from N (0 , A − 1 ) with an unknown precision matrix A ∈ R d
′ × d ′ , and for each i ≤ n , η i ∼ N (0 , B − 1 i ) are additive noise vectors with precision matrices B
i ∈ R m × m relating the m output dimensions.
Our goal is to infer the posterior distribution for the parameters θ given our observations, under the setting of A of the form A = αI 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 Π , for a given choice of A , and updating A , until the pair ( A, Π) converge to a locally optimal setting.
This corresponds to an EM algorithm (Bishop, 2006).
Henceforth, we will use the following stacked notation: we write Y ∈ R nm for the concatenation of y 1 , . . .
, y n ; B ∈ R nm × nm for a block diagonal matrix with blocks B 1 , . . .
, B n and Φ =
[ ϕ ( X 1 ) T ; . . .
; ϕ ( X n ) T ] T ∈ R nm × d ′ for the embedded design matrix.
We write M :=
Φ T BΦ .
Additionally, for a vector v and a PSD matrix G of compatible dimensions, ∥ v ∥ 2 G = v T Gv .
With that, the aforementioned EM algorithm starts with some initial A ∈ R d ′ × d ′ , and iterates:
• (E step) Given A , the posterior for θ , denoted Π , is computed exactly as
[Equation]
where H = M + A and
[Equation]
• (M step) We lower bound the log-probability density of the observed data, i.e. the evidence, for the model with posterior Π and precision A ′ as (derivation in Appendix B.2)
[Equation]
for C independent of A ′ .
We choose an A that improves this lower bound.
Limited scalability
The above inference and hyperparameter selection procedure for Π and A is futile when both d ′ and nm are large.
The E-step requires the inversion of a d ′ × d ′ matrix and the M-step evaluating its log-determinant, both cubic operations in d ′ .
These may be rewritten to instead yield a cubic dependence on nm (as in Section 3.3), but under our assumptions, that too is not computationally tractable.
Instead, we now pursue a stochastic approximation to this EM-procedure.
E VIDENCE MAXIMISATION USING STOCHASTIC APPROXIMATION
We now present our main contribution, a stochastic approximation to MacKay’s iterative algorithm, presented in the previous section: an M-step requiring access only to samples from Π , and a method for obtaining approximate posterior samples through stochastic optimisation for the E-step.
3. H YPERPARAMETER SELECTION USING POSTERIOR SAMPLES (M- STEP )
For now, assume that we have an efficient method of obtaining samples ζ 1 , . . .
, ζ k ∼ Π 0 at each step, where Π 0 is a zero-mean version of Π , and access to ¯ θ , the mean of Π .
Evaluating the first order optimality condition for M (see Appendix B.3) yields that the optimal choice of A satisfies
[Equation]
where the quantity γ is the effective dimension of the regression problem.
It can be interpreted as the number of directions in which the weights θ are strongly determined by the data.
Setting A = αI for α = γ/ ∥
¯ θ ∥ 2 yields a contraction step converging towards the optimum of M (Mackay, 1992).
Computing γ directly requires the inversion of H , a cubic operation.
We instead rewrite γ as an expectation with respect to Π 0 using Hutchinson (1990)’s trick, and approximate it using samples as
[Equation]
We then select α = ˆ γ/ ∥
¯ θ ∥ 2 .
We have thus avoided the explicit cubic cost of computing the logdeterminant in the expression for M (given in (3)) or inverting H .
Due to the block structure of B , ˆ γ may be computed in order n vector-matrix products.
3. S AMPLING FROM THE LINEAR MODEL ’ S POSTERIOR USING SGD (E- STEP )
Now we turn to sampling from Π 0 = N (0 , H − 1 ) .
It is known (also, shown in Appendix C.1)
that for
E ∈ R nm the concatenation of ε 1 , . .
, ε n with ε i ∼ N (0 , B − 1 i ) and θ 0 ∼ N (0 , A − 1 ) , the minimiser of the following loss is a random variable ζ with distribution Π 0 :
[Equation] [Equation]
This is known as the sample-then-optimise method (Papandreou & Yuille, 2010).
We may thus obtain a posterior sample by optimising this quadratic loss for a given sample pair ( E , θ 0 ) .
Examining L :
• The first term is data dependent.
It corresponds to the scaled squared error in fitting E as a linear combination of Φ .
Its gradient requires stochastic approximation for large datasets.
• The second term
, a regulariser centred at θ 0 , does not depend on the data.
Its gradient can thus be computed exactly at every optimisation step.
Predicting E , i.e. random noise, well from Φ is hard.
Due to this, the variance of a mini-batch estimate of the gradient of ∥ Φ z − E∥ 2B may be large.
Instead, for E and θ 0 defined as before, we propose the following alternative loss, equal to L up to an additive constant independent of z :
[Equation]
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 E and the prior sample θ 0 both feature within the regularisation term—the gradient of which can be computed exactly—rather than in the data-dependent term.
This may lead to a lower minibatch gradient variance.
To see this, consider the variance of the single-datapoint stochastic gradient estimators for both objectives’ data dependent terms.
At z ∈ R d ′ , for datapoint index j ∼ Unif ( { 1 , . . . , n } ) , these are
[Equation]
for L and L ′ , respectively.
Direct calculation, presented in Appendix C.2, shows that
[Equation]
Note that both Var ˆ g and Var ˆ g ′ 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 ˆ g ; that is if Tr ∆ > 0 .
We analyse two key settings:
• At initialisation, taking z = θ 0 (or any other initialisation independent of E ),
[Equation]
where we used that E [ EE T ]
= B − 1 and since E is zero mean and independent of θ 0 , we have E [ θ 0 E T ] =
E θ 0 E E T = 0 .
Thus, the new objective L ′ is always preferred at initialisation.
• At convergence, that is,
when z
= ζ , a more involved calculation presented in Appendix C.3 shows that L ′ is preferred if
[Equation]
This is satisfied if α is large relative to the eigenvalues of M (see Appendix C.4), so when the parameters are not strongly determined by the data relative to the prior.
When L ′ 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 L ′ is not preferred at convergence, it may still be preferred for most of the optimisation, before the noise is fit well enough.
3. D UAL FORM OF E- STEP : M ATHERON ’ S RULE AS OPTIMISATION
The minimiser of L and L ′ is ζ
= H − 1 ( Aθ 0 + Φ T B E ) .
The dual (kernelised) form of this is
[Equation]
which is known in the literature as Matheron’s rule (Wilson et al., 2020, our Appendix D).
Evaluating (12) requires solving a mn -dimensional quadratic optimisation problem, which may be preferable to the primal problem for mn < d ′ ; however, this dual form cannot be minibatched over observations.
For small n , we can solve this optimisation problem using iterative full-batch quadratic optimisation algorithms (e.g. conjugate gradients), significantly accelerating our sample-based EM iteration.
NN UNCERTAINTY QUANTIFICATION AS LINEAR MODEL INFERENCE
Consider the problem of m -output prediction.
Suppose that we have trained a neural network of the form f : R d ′ × R d (cid:55)→ R m , obtaining weights ¯ w ∈ R d ′ , using a loss of the form
[Equation] [Equation]
where ℓ is a data fit term (a negative log-likelihood) and R is a regulariser.
We now show how to use linearised Laplace to quantify uncertainty in the network predictions f ( ¯ w, · ) .
We then present the g-prior, a feature normalisation that resolves a certain pathology in the linearised Laplace method when the network f contains normalisation layers.
4. T HE LINEARISED L APLACE 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-Gaussian):
We take a first-order Taylor expansion of f around ¯ w , yielding the surrogate model
[Equation]
+
[Equation]
θ
[Equation]
for
[Equation] [Equation]
This is an affine model in the features ϕ ( x ) given by the network Jacobian at x .
We approximate the loss of the linear model
L ( h ( θ, · )) with a quadratic, and treat it as a negative log-density for the parameters θ , yielding a Gaussian posterior of the form
N (¯ θ, ( ∇ 2 θ L ) − 1 ( h (¯ θ, · )))
where
[Equation]
Direct calculation shows that ( ∇ 2 θ L )( h (¯ θ, · )) =
( A + Φ T BΦ) =
H , for ∇ 2 w R ( ¯ w )
=
A and
B a block diagonal matrix with blocks B i = ∇ 2ˆ y i ℓ ( y i , ˆ y i ) evaluated at ˆ y i
= h (¯ θ, x i ) .
We have thus recovered a conjugate Gaussian multi-output linear model.
We treat A as a learnable parameter thereafter.
In practice, we depart from the above procedure in two ways:
• We use the neural network output f ( ¯ w, · ) as the predictive mean, rather than the surrogate model mean h (¯ θ, · ) .
Nonetheless, we still need to compute ¯ θ to use within the M-step of the EM procedure.
To do this, we minimise L ( h ( θ, · )) over θ ∈ R d ′ .
• We compute the loss curvature B i at predictions ˆ y i =
f ( ¯ w, x i ) in place of h (¯ θ, x i ) , 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. O N THE COMPUTATIONAL ADVANTAGE OF SAMPLE - BASED PREDICTIONS
The linearised Laplace predictive posterior at an input x is N ( f ( ¯ w, x ) , ϕ
( x ) H − 1 ϕ ( x ) T ) .
Even given
H − 1 , evaluating this naïvely requires instantiating ϕ ( x ) , at a cost of m vector-Jacobian products (i.e. backward passes).
This is prohibitive for large m .
However, expectations of any function r : R m (cid:55)→ R under the predictive posterior can be approximated using only samples from Π 0 as
E [ r ] ≈ 1 k (cid:80) kj =1 r (
[Equation]
j ) for ψ j
[Equation]
ζ
[Equation]
with ζ 1 , . .
[Equation]
(16) requiring only k Jacobian-vector products.
In practice, we find k much smaller than m suffices.
4. F EATURE EMBEDDING NORMALISATION : THE DIAGONAL G - PRIOR
Due to symmetries and indeterminacies in neural networks, the embedding function ϕ ( · ) = ∇ w f ( ¯ w, · ) used in the 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 an isotropic prior; that is, considering ϕ ( x ) θ 0 for θ 0 ∼ N (0 , α − 1 I ) .
There are two natural solutions: either normalise the features by their (empirical) second moment, using the normalised embedding function ϕ ′ given by
[Equation]
given by s i =
[Φ T BΦ]
[Equation] [Equation]
(17) or likewise scale the prior, setting A
= α diag (( s − 2 i )
d ′ i =1 ) .
The latter formulation is a diagonal version of what is known as the g-prior (Zellner, 1986).
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 et al.
(2022c) show that normalisation layers lead to indeterminacies in 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 we show in Appendix E, the g-prior cancels these indeterminacies, allowing for the use of a single regularisation parameter.
D EMONSTRATION : SAMPLE - BASED LINEARISED L APLACE INFERENCE
We demonstrate our linear model inference and hyperparameter selection approach on the problem of 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 our 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 us to evaluate the quality of our approximations.
We then demonstrate our method at scale on CIFAR100 classification with a ResNet18 (Section 5.2), and the dual 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 − 1 , computing their logdeterminants, or instantiating Jacobian matrices ϕ ( x ) , all of which have hindered previous
linearised
Laplace implementations.
We interact with NN Jacobians only through Jacobian-vector and vectorJacobian products, which have the same asymptotic computational and memory costs as a NN forward-pass (Novak et al., 2022).
Unless otherwise specified, we use the diagonal g-prior, resulting in a scalar regularisation parameter.
We present a summary of our method in Algorithm 1, an illustrative example in Figure 1, and full algorithmic detail in Appendix F.
An implementation of our method in JAX can be found in an anonymised repository at this link .
Additional results on predictive distribution calibration are provided in Appendix H and Appendix I.
Algorithm 1: Sampling- and
Inputs: initial α > 0 ; k, k ′ ∈ N , number of samples for stochastic EM and prediction, respectively.
Compute g-prior scaling vector s as in (17)
Sample random regularisers θ n 1 , . .
, θ nk per (7) while α has not converged do
Find posterior mode ¯ θ by optimising linear model loss L ( h ( θ, · )) , given in (13)
Draw posterior samples ζ 1 . . .
ζ k by optimising objective L ′ with θ n 1 , . .
, θ nk
Estimate effective dimension ˆ γ , per (5), using samples ζ 1 . . .
ζ k
Update prior precision α ← ˆ γ/ ∥ ¯ θ ∥
Sample k ′ random regularisers θ n ′ 1 , . .
, θ n ′ k ′ using optimised α
Draw corresponding posterior samples ζ ′ 1 , . .
, ζ ′ k ′ using loss L ′
Output: posterior samples ζ ′ 1 , . . .
, ζ ′ k
[Figure]
prior
[Figure]
posterior functions sample optimisation after EM
α = 0.
° 1.5 ° 1.0 ° 0.5 0.0 0.5 1.0 1.5 x °°   y ° 1.5 ° 1.0 ° 0.5 0.0 0.5 1.0 1.5 x °°   y ° 1.5 ° 1.0 ° 0.5 0.0 0.5 1.0 1.5 x °°   y ° 1.5 ° 1.0 ° 0.5 0.0 0.5 1.0 1.5 x °°   y ° 20 q°  qpriorsamplesposteriorsamples ° 4 ° 2 4 q 0 ° 5.
° 2.
q
5. A BLATION STUDY : L E N ET 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 =60 k observations and we employ
LeNet-style CNNs of increasing size: “LeNetSmall” ( d ′ =14634 ), “LeNet” ( d ′ =29226 ) and “LeNetBig” ( 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 in the literature L , using LeNet.
The results are shown in Figure 2.
We draw exact samples: ( ζ ⋆j ) j ≤ k through matrix inversion and assess sample fidelity in terms of normalised squared distance to these exact samples ∥ ζ − ζ ⋆ ∥ 22 / ∥ ζ ⋆ ∥ 22 .
All runs share a prior precision of α ≈ 5 . 5 obtained with EM iteration.
The effective dimension is ˆ γ ≈ 1300 .
Noting that the g-prior feature normalisation results in Tr M = d ′ , we can see that condition (11) is not satisfied (2 × 5 .
5 × 1300 < 29 k ) .
Despite this, the proposed objective converges to more accurate samples even when using a 16-times smaller batch size (left plot).
The right side plots relate sample error to categorical symmetrised-KL (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 a sample error of ≈ 0 .
5 but, nevertheless, the proposed loss L ′ reaches lower a prediction error.
Fidelity of sampling inference We benchmark our method 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 full-covariance lin.
Laplace predictive posterior.
As standard, we compute categorical predictive distributions with the probit approximation (Daxberger et al., 2021a).
All methods use the same layerwise prior precision obtained with 5 steps of full-covariance EM iteration.
For all three LeNet sizes, the sampled approximation presents the lowest categorical sym.
[Figure] [Figure]
difference with full lin.
Laplace predictive
[Figure]
KL and logit W2 distance to the exact lin.
Laplace pred.
posterior (Figure 3, LHS).
Surprisingly, while the fidelity of competing approximations degrades with model size, that of sampling increases.
Accuracy of sampling hyperparameter selection
We compare our sampling EM iteration with 16 samples to the 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 to the same value.
In this setting, the diagonal covariance approximation diverges, and KFAC converges to a 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, are needed to accurately predict the targets; other features are pruned away.
5. R ES N ET 18 ON CIFAR
We showcase the scalability of our approach by applying it to CIFAR100 m =100 -way image classification.
The training set consists of n =50 k observations, and we employ a ResNet-18 model with d ′ ≈ 11 M parameters.
We perform optimisation using SGD with Nesterov momentum and a linear learning rate decay schedule.
Unless specified otherwise, we run 8 steps of EM with 6 samples to select α .
We then optimise 64 samples to be used for prediction.
We run each experiment withdifferent seeds reporting mean and std.
error.
Full experimental details are in Appendix G.
Stability and cost of sampling algorithm
Figure 4 shows that our sample-based EM converges in 6 steps, even when using a single sample.
At convergence, α ≈ 10 4 and ˆ γ ≈ 700 , so 2 × 700 × 10 4 = 1 . 4 × 10 7 > 1 . 1 × 10 7 .
Thus, (11) is satisfied and L ′ 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 Ovadia et al.
(2019) benchmark for evaluating methods’ test Log-Likelihood (LL) on the increasingly corrupted data sets of Hendrycks & Gimpel (2017).
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 approximation of the lin.
Laplace covariance (Ritter et al., 2018).
For the
[Figure] [Figure]
[Table]
[Figure] [Figure]
[Table]
latter, constructing full Jacobian matrices for every test point is computationally intractable, so we use 64 samples for prediction, as suggested in Section 4.2.
The KFAC covariance structure leads to fast log-determinant computation, allowing us to learn layer-wise prior precisions (following Immer et al., 2021a) 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 et al., 2021b).
Figure 5 shows that for in-distribution inputs, ensembles performs best and KFAC overestimates uncertainty, degrading LL relative to point-estimated 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.
We group our test-set into sets of κ data points and then uniformly re-sample the points in each set until sets contain τ points.
We then evaluate the LL of each set jointly.
Since each set only contains κ distinct points, a predictor that models self-covariances perfectly should obtain an LL value at least as large as its marginal LL for all values of κ .
We use τ =1000 and repeat the experiment for 10 test-set shuffles.
Our setup remains the same as above, except that we use Monte Carlo marginalisation instead of probit, since the latter discards covariance information.
Table 1 shows that ensembles make calibrated predictions marginally but their covariance estimates are poor, an observation also made by Osband et al.
Our approach is competitive for all κ , performing best in the challenging large κ cases.
5. T OMOGRAPHIC RECONSTRUCTION
To demonstrate our approach in dual (kernelised) form, we replicate the setting of Barbano et al.
(2022) and Antorán et al.
(2022b), where linearised Laplace is used to estimate uncertainty for a tomographic reconstruction output 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 more details.
Whereas Antorán et al.
(2022b) use a single step of exact EM to learn hyperparameters, we use our sample-based variant and runsteps.
Unless otherwise specified, we use 16 samples for stochastic EM, and 1024 for prediction.
We test on the real-measured µ CT dataset of 251 k pixel scans of a single walnut released by
Der Sarkissian et al.
(2019b).
We train d ′ =2 .
97 M parameter U-Nets on the m =7680 dimensional observation used by Antorán et al.
(2022b) and a twice as large setting m =15360 .
Here, the U-Net’s input is clamped to a constant ( n = 1 ), and its parameters are optimised to output the reconstructed image.
As a result, we do not need mini-batching and can draw samples using Matheron’s rule (12)
[Figure]
We solve the linear system contained therein using conjugate gradient (CG) iteration implemented
[Figure]
[Table]
with GPyTorch (Gardner et al., 2018).
We accelerate CG with a randomised SVD preconditioner of rank 400 (alg.
5.6 in Halko et al., 2011).
See Appendix G for full experimental details.
Stability and cost of sampling algorithm Figure 6 shows that sample-based EM iteration converges within 4 steps using as few as 2 samples.
Table 2 shows the time taken to perform 5 sample-based EM steps for both the m =7680 and m =15360 settings; avoiding explicit estimation of the covariance log-determinant provides us with a two order of magnitude speedup relative to Antorán et al. (2022c) for hyperparameter learning.
By avoiding covariance inversion, we obtain an order of magnitude speedup for prediction.
Furthermore, while scaling to double the observations m =15360 is intractable with the previous method, our sampling method requires only a 25% increase in computation time.
Predictive performance Figure 7 shows qualitatively that the marginal standard deviation assigned to each pixel by our method aligns with the pixelwise error in the U-Net reconstruction in a finegrained 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 more uniformly across large sections of the image.
Table 2 shows that the pixelwise LL obtained with our method exceeds that obtained by Antorán et al.
(2022b), potentially due to us optimising the prior precision to convergence while the previous work could only afford a single EM step.
The rightmost plot in Figure 6 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.
[Figure]
C ONCLUSION
Our work introduced a sample-based approximation to inference and hyperparameter selection in
Gaussian linear multi-output models.
This allowed us to scale the linearised Laplace method to
ResNet18 on CIFAR100, where it was computationally intractable with existing methods.
We also demonstrated the strength of the approach on a high resolution tomographic reconstruction task, 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 are of importance, and computation of posterior samples is often needed.
R EPRODUCIBILITY S TATEMENT
In order to aid the reproduction of our results, we provide a high-level overview of our procedure in algorithm 1 and the fully detailed algorithms we use in our two major experiments in Appendix F.
Appendix G provides full experimental details for all datasets and models used in our experiments.
Our code is available in an anonymised repository at this link .
A E XTENDED 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 a 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 et 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.e. kernelised or Gaussian Process) form.
Additionally, for non-Gaussian likelihoods, e.g. in classification, inference is no longer closed form.
The most common approximations used in these settings are Laplace’s method (Mackay, 1992) and variational inference (Hensman et al., 2013).
Khan et al.
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-optimise to Laplace posteriors.
Sample-then-optimise Papandreou & Yuille (2010); de G. Matthews et al.
(2017) phrase sampling from a conjugate Gaussian-linear model as solving a perturbed quadratic optimisation problem.
This method has been applied for uncertainty estimation in non-linearised NNs by Osband et al.
(2018; 2021), and Pearce et al.
(2020), although in this setting it does not draw exact posterior samples.
In this work, we show sample-then-optimise to be the primal form of Matheron’s rule (Journel &
Huijbregts, 1978; Hoffman & Ribak, 1991), a method for updating jointly Gaussian samples into conditional samples, which was recently repopularised by Wilson et al.
Linearised neural networks Introduced by Mackay (1992), these are an approximation yielding 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.
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 point-estimated NNs’ predictions.
Immer et al.
(2021a)
and Antorán et al.
(2022a;c) explore the linear model’s evidence for model selection.
Daxberger et al.
(2021b)
and
Maddox et al.
(2021) introduce subnetwork and finite differences approaches, respectively, for faster inference with the linearised model.
This line of work is intimately related to the neural tangent kernel (Jacot et al., 2018; Lee et 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, 2011; Baragatti & Pommeret, 2012).
In the large-scale setting, we overcome the computational intractability of the Fisher by diagonalising the g-prior while preserving its invariance property.
B M ODEL EVIDENCE LOWER BOUND AND THE EFFECTIVE DIMENSION B. E QUIVALENT FORMULATIONS OF EFFECTIVE DIMENSION
We begin by relating two standard forms of effective dimension, which we use throughout.
Starting with the form standard in the kernel-based literature (that without an explicit d ′ dependence),
[Equation]
we have arrived at the form used within the finite-dimensional linear modelling literature (Mackay, 1992; Wipf & Nagarajan, 2007; Maddox et al., 2020).
B. D ERIVATION OF M AS A LOWER BOUND ON THE MODEL EVIDENCE
Let p θ be the Lebesgue density of N (Φ θ, B − 1 ) , P = N (0 , A ′− 1 ) and Q = N (¯ θ, ( M + A ′ ) − 1 ) .
Then,
[Equation]
where D denotes the KL-divergence.
Starting with the first term,
[Equation]
and expanding the quadratic form,
[Equation]
[Equation]
To handle the final integral
[Equation]
consider that
[Equation] [Equation]
and thus
[Equation]
The KL between two multivariate Gaussians is a standard result, yielding
[Equation]
where we used that γ = d ′ − Tr { A ′ ( M + A ′ ) − 1 } .
Putting together (31) and (33), we obtain
[Equation]
[Equation]
which is the stated result up to taking C = log p ¯ θ ( Y ) .
B. F IRST ORDER OPTIMALITY CONDITION FOR M
Consider the derivative of M .
We have,
[Equation] [Equation]
where we expanded log det ( I + A − 1 M )
= log det ( A + M ) − log det A .
Taking the respective derivatives and setting equal to zero at A , this leads to the condition
[Equation]
Post-multiplying by A and applying the push-through identity, we obtain
[Equation]
For the above to hold, it is necessary that the traces of both sides are equal.
Thus,
[Equation]
[Equation]
which is the stated first order optimality condition, up to a cyclic permutation.
B. M- STEP FOR FEATURE - WISE REGULARISATION STRENGTHS
We can leverage the primal form expression for the effective dimension given in Appendix B.1 to extend the above first order optimality condition to the feature-wise regulariser setting.
Consider a sub-vector of our weight vector contiguous between the i th and j th weights written as ¯ θ i : j .
Note that we only choose contiguous weights for notational convenience but it is not necessary to do so in general.
The first order condition from Appendix B.3 is satisfied if for any i, j with i <
j we have
[Equation]
and thus we may update the regulariser for each separate sub-vector as α = γ
i : j / ∥
¯ θ i : j ∥ 2 .
C A NALYSIS OF LOSSES AND LOSS GRADIENT ESTIMATOR VARIANCES C. O N LOSS MINIMA
The losses L and L ′ are strictly convex, thus to confirm they have the same unique minimum, it suffices to consider the respective first order optimality conditions, ∇ L ( ζ ) = 0 and ∇ L ′ ( ζ ′ )
= 0 .
We have,
[Equation] [Equation]
and
[Equation] [Equation]
Thus ζ = ζ ′ almost surely.
Moreover, L ′ ( z ) = L ( z ) + C for all z , for C a constant independent of z .
To determine the distribution of ζ , note that it is a linear transformation of zero-mean Gaussian random variables, and thus itself a zero-mean Gaussian random variable.
Rearranging the first order optimality condition, we find that
[Equation] [Equation]
Thus
[Equation] [Equation]
[Equation]
And so ζ ∼ N (0 , H − 1 ) =
Π 0 as claimed.
C. L OSS GRADIENT VARIANCE CONDITION
Taking j ∼ Unif ( { 1 , . .
. , n } )
[Equation]
the gradient estimators for the data-dependent terms of L and L ′ are
[Equation]
and
[Equation]
respectively.
Their variances are related as
[Equation]
Evaluating the variance and covariance
[Equation]
we have
[Equation]
and
[Equation]
and thus
[Equation]
C. C ONDITION AT CONVERGENCE
Now consider Tr ∆ for z = ζ ∼ Π 0 , the optimum of both L and L ′ .
From the first order optimality condition,
[Equation]
Proceeding to rearrange the condition at z = ζ ,
[Equation]
where we substituted in the definition of ζ , then used that E and θ 0 are independent, and that E [ θ 0 ] = 0 , and finally that E [ EE T ]
= B − 1 .
Writing M =
Φ T BΦ and recalling that H
=
( M + A ) , we have
[Equation] [Equation]
where we have used that ( A − 1 M + I ) − 1 A − 1 M = I − ( A − 1 M
+ I ) − 1 for the fourth equality.
Now consider the isotropic prior case A = αI and recall the effective dimension is written as γ =
Tr { M ( M + A ) − 1 } .
The above
C. A NALYSING CONDITION AT CONVERGENCE
To gain some intuition for the condition at convergence, denote by λ 1 , . .
λ d ′ the eigenvalues of M (with multiplicity).
We can use these to restate the condition as
[Equation]
This formulation of effective dimension gives an interpretation of a soft count of the number of dimensions for which λ j is larger than α ; in that sense, λ j measures how well determined the corresponding dimension of the weight vector θ is by the observed data.
From here, note that
[Equation]
and thus it is sufficient for Tr ∆ > 0 to hold at convergence that α > λ j for all j (but, of course, not necessary), yielding the intuition that L ′ is preferred when the problem is heavily regularised.
D D UAL FORM OF THE SAMPLE - THEN - OPTIMISE LOSS : M ATHERON ’ S RULE
Both losses L and L ′ result in a random variable ζ ∼ Π 0 given by
[Equation]
Recalling that H
=
A +Φ T BΦ and using the push-through identity, we can express ζ equivalently as
[Equation]
[Equation] [Equation]
[Equation] [Equation]
Now taking a sample of the posterior Gaussian process evaluated at input x to be G
= ϕ ( x ) ζ and the corresponding sample of the prior process to be G
= ϕ
( x ) θ 0 , premultiplying the above expression by ϕ ( x ) we obtain
[Equation]
which is Matheron’s rule.
E R ESOLVING FEATURE SCALE INDETERMINACIES IN THE NN J ACOBIAN WITH THE G - PRIOR
Antorán et al.
(2022c) show that for NNs with normalisation layers, the Jacobian features ϕ ( · )
= ∇ w f ( ¯ 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 ¯ w =
[ ¯ w 0 , ¯ w 1 ] .
We assume the layer containing ¯ w 0 is followed by a normalisation layer, but not that containing ¯ w 1 , which leads to the invariance
[Equation]
[Equation]
for all k > 0 .
While f is invariant to this scaling, the Jacobian feature embeddings ϕ ( · ) = ∇ w f ( ¯ w, · ) are not.
We separate the embeddings as
[Equation]
Antorán et al.
(2022c) show that, given a reference pair ([ ¯ w 0 , ¯ w 1 ] , [ ϕ 0 ( x ) , ϕ 1 ( x )]) , and for ¯ wnormalised, scaling ¯ w 0 by k results in the pair ([ k ¯ w 0 , ¯ w 1 ] , [ k − 1 ϕ 0 ( x ) , ϕ 1 ( x )]) .
Thus, using a single prior precision parameter, the regularisation strength applied to the weights multiplying ϕ 0 ( x ) relative to those multiplying ϕ 1 ( x ) will increase with k .
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 to resolve this is to assign the weights ¯ w 0 and ¯ w 1 separate regularisation parameters and learn these using the EM procedure outlined in Section 2.
However, instead, consider using the g-prior normalised features ϕ ′ introduced in Section 4.3, and specifically, the scaling vector corresponding to normalised and non-normalised components s =
[ s 0 , s 1 ] .
For a reference pair ([ ¯ w 0 , ¯ w 1 ] , [ s 0 , s 1 ]) and for ¯ w 0 normalised, the k-scaled pair is ([ k ¯ w 0 , ¯ w 1 ] and
[Equation]
where ⊙ denotes an elementwise power.
Since the k -scaled features are [ k − 1 ϕ 0 ( · ) , ϕ 1 ( · )] , when applying the g-prior normalisation we recover a 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 1 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 we use for tomographic image reconstruction.
Image classification Algorithm 2 provides an algorithm for linearised Laplace inference using the stochastic EM 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 x i , denoted B i , is given by B i =
diag ( p i ) − p
i
p
Ti for p i = µ ( f ( ¯ w, x i ))
our neural network’s 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 a positive definite matrix as its 1 / 2 th power.
In order to limit computational cost, we sample the stochastic regularisation terms ( θ nj ) , per (7), only once at the start.
Not resampling these at each E step results in the optima of the sampling objective being close for successive iterations.
This comes at the cost of a small bias in our estimator which we find to be negligible in practise.
We separate ( θ nj ) into a sum consisting of a prior sample from ( θ 0 j ) and a data dependent term, denoted ( θ ′ j ) .
The former scales with α − 1 / 2 while the latter with α − 1 so this allows us to update each term in closed form each time α changes.
We initialise our samples at ( θ 0 j ) 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 θ ′ 1 as s = α − 1 ( E [ θ ′ 1 ⊙ θ ′ 1 ]) ⊙− 1 / 2 .
We optimise both our samples ζ
and posterior mean ¯ θ using stochastic gradient descent with a
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 L is non-quadratic and thus Polyak averaging is 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 hyperparameters of our algorithm are the number of samples to draw for the EM iteration, the number of EM steps to run, and SGD hyperparameters: learning rate, linear decay rate, number of steps, batch-size and gradient clipping.
Empirically, we find that at most 5 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.
Inputs: Linearised network h , linearisation point ¯ w , observations x 1 , . . .
, x n , negative log-likelihood function ℓ , initial precision α > 0 , number of samples k
Function B(i) : p i ← µ ( h ( ¯ w, x i )) return diag ( p i ) − p
i p Ti for j = 1 , . . .
, k do θ
j ∼ N (0 , α − 1 I ) θ ′
j ← α − 1 (cid:80) ni =1 ϕ ( x i ) T ε j where ε j ∼ N (0 , B ( i ))
ζ j ← θ
j ¯ θ ←s ← α − 1 (cid:104) 1 k (cid:80) kj =1 θ
′⊙ 2 j (cid:105) ⊙− 1 /while
α has not converged do for j = 1 , . . .
, k do ζ j ← SGD z (cid:0) ∥
Φ( s ⊙ z )
∥ 2B + α ∥
z − θ
j − ( s ⊙ θ ′ j ) ∥ 22 , init =
ζ j (cid:1) ¯ θ ← SGD θ (cid:0)(cid:80) ni =1 ℓ ( y i , h (( s ⊙ θ ) , x i ) + α ∥ θ ∥ 22 , init =
¯ θ (cid:1) ˆ γ ← 1 k (cid:80) k
j =1 (cid:80) n i =1 ∥ ( ζ j ⊙ s ) T ϕ ( x i ) T B ( i ) 12 ∥ 2α ′ ← ˆ γ/ ∥
¯ θ ∥ 22 for j = 1 , . .
, k do θ
j ← (cid:112) αα ′ θ
j θ ′ j ← αα ′ θ ′
j α ← α ′ Output: Optimised precision α and weight samples ζ 1 , . . .
, ζ k
Inputs: Linearised network h , linearisation point ¯ w , measurements Y , discrete Radon transform
U , U-Net Jacobian Φ , initial precision α > 0 , number of samples k , noise precision B
Function Kvp( v , α , U Φ , s, B − 1 ) : return U Φ( α − 1 diag ( s ⊙ 2 ))
Φ T U T v + B − 1 v s ← ( (cid:80) i<m ( U i Φ) ⊙ 2 ) − 1 /while
α has not converged do
P ← Compute-preconditioner ( Kvp ) for j = 1 , . .
, k do ζ
j ← U Φ( s ⊙ θ 0 j ) + E j where E j ∼ N (0 , B − 1 ) and θ
j ∼ N (0 , A − 1 ) c
j ← CG (cid:0) Kvp , ζ 0 j , precond.
= P (cid:1) ζ
j ← ζ 0 j − U Φ( α − 1 diag ( s ⊙ 2 ))
Φ T U T c j δ ← U
(Φ ¯ w − f ( ¯ w ))
c ← CG ( Kvp , Y + δ, precond.
= P ) ¯ θ ← s ⊙ α − 1 Φ T U T c
ˆ γ ← 1 k (cid:80) kj =1 ∥
U Φ( s ⊙ ζ j ) ∥α ′ ← ˆ γ/ ∥
¯ θ ∥ 2 α ← α ′ Output: Optimised precision α
As a result, we can tune these parameters on the train data, no validation set is 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 our classification networks, the constant-in-
θ terms cancel in the linearised model and we are left with h ( θ, x )
= ϕ
( x ) θ (Antorán et al., 2022c).
In our algorithm, this fact is only relevant to the computation of the posterior mode ¯ θ as the optima of L ( h ( θ, · )) .
Tomographic reconstruction Algorithm 3 is the kernelised version of algorithm 2 that we use for tomographic reconstruction.
This problem is described in detail in Appendix G.3.
Distinctly from the image classification setting, tomographic reconstruction is a regression problem for which we use a Gaussian likelihood with fixed noise precision B = I .
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 loss can be optimised in closed form by solving a linear system of equations given by the observation covariance, i.e. the kernel matrix, U Φ( α − 1 diag ( s ⊙ 2 ))
Φ T U T + B − 1 where the linear operator U represents the discrete Radon transform and combines with the U-Net Jacobian to build the feature embedding U Φ .
We solve against the kernel matrix using the preconditioned conjugate gradient (CG) method described by Gardner et al.
As a preconditioner, we compute a 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 ( θ 0 , E ) at 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 α .
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 use of a large preconditioner results in CG always hitting the desired low error tolerance within 10 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.
A particularity of this setting is that the U-Net does not have a dense final layer.
As a result, the constant-in- θ terms in the linearised function h do not cancel (see Section 4.1), leading to the appearance of the target offset term δ when solving for the posterior mean.
G E XPERIMENTAL DETAILS
In this appendix we provide experimental details and hyperparameter settings omitted from the main text.
G. MNIST EXPERIMENTS
MNIST m =10 way classification experiments were performed using the LeNet-style CNN architectures of increasing size employed by Antorán et al.
(2022c): “LeNetSmall” ( d ′ =14634 ), “LeNet” ( d ′ =29226 ) and “LeNetBig” ( d ′ =46024 ).
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 a learning rate drop of a factor of 10 every 30 epochs.
The MNIST dataset was downloaded from PyTorch torchvision .
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 our SGD approaches to convex optimisation for obtaining posterior modes and samples are as follows
• 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 a maximum norm of 1 .
We use an initial learning rate of 1 e − 2 when using standard isotropic or layerwise Gaussian priors, and 1 for the g-prior.
We employ a linear decay schedule that reduces the lr by a factor of 330 over the first 75 % of the training procedure and holds it constant afterwards.
Sampling :
We 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 use a learning rate of 2 e − 1 , whereas for the g-prior, we find a higher learning rate of 200 to work best.
Hyperparameter optimisation
: We tuned the learning rate, decay schedule and gradient clipping strength using a 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 of learning a single precision hyperparameter and layerwise hyperparameters in Figure 3, we extend the M-step update to as α l = γ l / ∥ ¯ θ l ∥ 22 where l indexes each layer’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 to the covariance is constructed by first computing the diagonal of the Hessian M and the inverting it.
For the KFAC covariance approximation, we exploit the equivalency between the Generalised Gauss Newton matrix (i.e. 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 a single sample per training observation, as in (Daxberger et al., 2021a).
For completeness, we also state the probit approximation for sampled predictive posteriors over logits.
For input x and samples ζ i , . . .
, ζ k , the predictive probability for class i ∈ ∥ 1 , . . .
, m ∥ is given by
[Equation]
G. CIFAR100 CLASSIFICATION
CIFAR100 m =100 way classification experiments were performed using ResNet18 models ( d ′ ≈ 11 M ) with specific architecture details matching the PyTorch torchvision implementation .
We train these models using SGD with momentum of 0.9 for 300 epochs.
The starting lr is 0.1 and we reduce it by a factor of 10 every 100 epochs.
The CIFAR100 dataset was also downloaded using torchvision 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 a batch size of 2000 for 40 epochs.
We employ a linear decay learning rate schedule with an initial learning rate of 1 e − 1 .
It is decreased by a 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 6 samples in parallel using SGD with Nesterov momentum (=0 . 9) and a batch size 100 for 20 epochs.
All other details match those of posterior mode optimisation.
Upon convergence of 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 batch size that for which we could simultaneously optimise 6 samples in parallel on a single hardware accelerator.
Similarly to the MNIST experiments, we did not optimise the number of optimisation epochs and instead chose large values that would ensure convergence.
It is likely that our EM iteration can be sped up by decreasing the duration of the convex optimisation routines.
Details for baselines and hyperparameters not mentioned explicitly in this subsection match those given for MNIST in the previous subsection.
G.2. E FFICIENT κ - ADIC SAMPLING
Osband et al.
(2022; 2021) introduced dyadic test input sampling ( κ = 2) as a practical way of evaluating joint predictions in discriminative tasks.
This method samples κ = 2 random anchor points from the test dataset, and then randomly resamples them to create a batch of size τ = 10 .
Test log-likelihood is evaluated jointly for each batch as
[Equation]
for f the model being evaluated and Π its posterior distribution over model parameters.
This quantity can be estimated with posterior samples ζ 1 , . . .
, ζ k ∼ Π as
[Equation]
We extend this evaluation approach to larger κ and τ values without increasing computational cost.
We randomly sample κ integers { b 1 , . . .
, b κ } such that they sum to τ , i.e (cid:80)
κi k i =
τ .
The joint log-likelihood over the batch of size τ with κ unique datapoints can then be estimated as
[Equation]
where the inner sum is over the κ distinct elements in the batch instead of the “total batch size” τ .
This is equivalent to the formulation proposed in Osband et al.
(2022) for dyadic sampling, when κ =and τ = 10 .
We note that it is not possible to achieve augmented dyadic sampling, as described in (Osband et 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 to 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 in our results from the choice of the κ observations which get grouped together in each batch.
G. T OMOGRAPHIC RECONSTRUCTION
Problem setup Tomographic reconstruction consists in solving a linear inverse problem in imaging where we observe a set of measurements y ∈ R m , which we assume to be generated as y =
Ux ∗ + η for U ∈ R m ×
d the discrete Radon transform, x ∗ ∈ R
d
the image to reconstruct and η ∼ N (0 , I ) random noise.
We have m ≪ d , making the problem underconstrained.
Our specific tomographic reconstruction task closely follows the one from Barbano et al.
We perform a sparse-view reconstruction of an image of a slice of a walnut from a sub-sampled set of measurements.
Specifically, from the full measurement set of (Der Sarkissian et al., 2019a), which containing scans atequidistant 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.
; Antorán et al. (2022b), we reduce the original 3D scan geometry to the 2D slice of 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 (501 px ) 2 and outputting a measurement tensor coherent with the described geometry.
[Figure]
posterior credible interval 0.
e m p
i
r
i c a l c o v e r a g e
Image reconstruction predicted error vs empirical error sampledlin.
U-Net lin.
U-Net
MCDOU-Net 0.
1.0 predicted probability 0.
e m p
i
r
i c a l a cc u r a c y
CIFAR100 predicted probability vs empirical accuracy
[Figure]
Methods To provide a reconstruction, we use the Deep Image prior (Ulyanov et al., 2020) which trains the parameters w ∈ R d ′ of a fully convolutional U-Net autoencoder f : R d ′ → R d , where the input is fixed and thus absent from our notation, until a satisfactory reconstruction f ( ¯ w ) 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 et al.
To estimate the uncertainty in this reconstruction, we linearise the U-Net around ¯ w , as described in Section 4.1.
This leaves us with a model affine in the parameters and with design matrix U Φ ∈ R m × d ′ .
We may now proceed with linear model inference.
While (Antorán et al., 2022b) use the traditional EM iteration described in Section 2, we use the sample-based one from Section 3.
For evaluation, we use the non-sparse reconstruction (using 1200 angles) provided by (Der Sarkissian et al., 2019a) as the ground truth image x ∗ .
To evaluate joint log-likelihoods 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): ˆΣ = 12 k (cid:104)(cid:80) kj =1 ˆ x 2 j + ˆ x j ˆ x Tj (cid:105) for (ˆ x j ) kj =1 samples from the predictive posterior over a patch.
We then construct predictive distributions over pixels as
N ( f ( ¯ w ) , ˆΣ) .
Hyperparameters We employ a low CG tolerance of 1 e − 3 and a maximum number of iterations of 150 , which is never reached in practise as the error tolerance level is always hit in less steps.
H C ALIBRATION OF PREDICTIVE DISTRIBUTIONS
This appendix evaluates the calibration of the predictive distributions provided by the methods under consideration in our CIFAR100 classification experiment (Section 5.2) and U-net image reconstruction experiment (Section 5.3).
For classification, we 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 at all confidence levels whereas MAP underestimates it.
Both sample-based 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 is the centered unit variance Gaussian.
We consider posterior credible intervals centered at 0 and of increasing width and plot the proportion of test points
[Table]
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, as in (Antorán et al., 2022b), consistently overestimates uncertainty.
Our approach, which performs 5 steps of EM, overestimates uncertainty
[Table]
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
[Table]
I A DDITIONAL E XPERIMENTS I. CIFAR100 C LASSIFICATION
Additional Baselines for marginal LL .
In the main text, we report the predictive uncertainty of our method versus point-estimated NNs (MAP), 5 ensembles of point-estimate NNs (Ensemble 5), and a KFAC estimate (KFAC).
Here, we report further comparisons with other baselines standard-inliterature: a diagonal approximation of the Laplace covariance matrix (diag), a Laplace approximation over a selected subnetwork of the original NN (subnetwork*), and a Laplace approximation over the last-layer of the NN further approximated by a KFAC matrix (KFAC-LL*).
We distinguish the last two methods in a separate category (*), since these methods require cross-validation with a held-out set in order to tune hyperparameters.
For these methods, we use 50 held-out points from the test set, and evaluate on the remaining 9950 points.
Figure 9 shows that somewhat surprisingly,
KFAC-LL performs quite competitively with our approach, and even with deep ensembles.
However, it is constrained to require a cross-validation set in order to tune hyperparameters.
Predictive uncertainty vs number of samples .
In the main text, we report predictive performance for our method with 64 samples.
However, it is possible that drawing many samples can be computationally expensive.
We compare the quality of predictive uncertainty obtained on the in-distribution test set and 5 corruptions vs the number of samples used for estimating the predictive variance, and results are shown in Figure 10.
Additional baselines for joint LL .
Similar to the marginal predictions, we also report KFAC-LL and subnetwork inference as two additional baselines for κ -adic sampling in Table 3.
KFAC-LL is once again quite competitive with our approach, but requires a held-out validation set.
I. T OMOGRAPHIC R ECONSTRUCTION
Measuring prediction calibration Figure 11 and Figure 12 are complementary to Figure 6 and
Figure 7, and report EM iteration and performance for m =15360 .
In Figure 13 and Figure 14, we show additional insights on the calibration of the uncertainty estimates obtained via the proposed method when compared to the MCDO approach.
Corruption severity −−−−−t e s t
LL
KFAC
MAP
MC Dropout
Ensemble (5) diag
KFAC-LL * subnetwork*
Sampling-MC
Sampling-Probit
4 83664 − 1.
− 1. − 1.
t e s t
L L
Split4 836 − 2.
− 2.
− 1.
t e s t
L L
Split4 836 − 3.
− 2.
− 2. − 2.
t e s t
L L
Split4 836 samples − 3.
− 3.
− 2.
t e s t
L L
Split4 836 samples − 4.
− 4.
− 3. − 3.
t e s t
L L
Split4 836 samples −−t
e s t
L L
Split
Sampling-Probit Sampling-MC
[Figure]
[Figure]
[Figure] [Figure]