Linearised Laplace Inference in Networks with Normalisation Layers and the Neural g-Prior
Abstract
We show that for neural networks (NN) with normalisation layers, i.e. batch norm, layer norm, or group norm, the Laplace model evidence does not approximate the volume of a posterior mode and is thus unsuitable for model selection.
We instead propose to use the
Laplace evidence of the linearized network, which is robust to the presence of these layers.
We also identify heterogeneity in the scale of Jacobian entries corresponding to different weights.
We ameliorate this issue by extending the scale-invariant g-prior to NNs.
We demonstrate these methods on toy regression, and image classification with a CNN.
1. Introduction
Normalisation layers, for example batch norm (Ioffe and Szegedy, 2015) or layer norm (Ba et al., 2016), are ubiquitous in modern NN architectures (He et al., 2016; Vaswani et al., 2017).
They speed up training and make it robust to the choice of hyperparameters.
They also introduce an invariance to the scale of the weights.
We study the effects of this invariance on probabilistic inference, in particular in the context of the linearised Laplace approximation, introduced by Mackay (1992).
This method approximates a NN with a surrogate linear model where the NN’s Jacobian acts as a feature expansion (Khan et al., 2019).
The linearised model shares its predictive mean with the NN.
Linearised Laplace predictive uncertainty is simple to compute (Immer et al., 2021b) and performs strongly on uncertainty quantification tasks (Kristiadi et al., 2020; Daxberger et al., 2021b,a).
The Laplace model evidence or marginal log-likelihood (MLL) can also be computed in closed form (Mackay, 1992) providing an objective for model selection.
Daxberger et al.
(2021b) find that the quality of uncertainty estimates provided by linearised Laplace is heavily dependent on the choice of Gaussian prior precision. They −  2 x −−   y
MLL with NN weights w ? −  2 x y
MLL with lin weights v ?
−  2 x y g-prior MLL with lin weights v ?
select an isotropic precision using cross validation.
The MLL is a preferable alternative, as it does not require a held out set and scales to more hyperparameters (Immer et al., 2021a).
We show that in normalised networks, the point around which the NN is linearised can not be a mode of the posterior for neither the NN model nor the linearised NN.
This results in the traditional formulation of the Laplace MLL returning spurious values, rendering it unsuitable for model selection.
We show how to find the mode of the linearised NN posterior and use it to propose a linearised Laplace MLL objective that is robust to the presence of normalisation layers.
Next, we observe heterogeneity in the scales of the NN Jacobian entries corresponding to different weights.
In this setting, isotropic priors can be excessively restrictive.
We address this by extending the scale invariant g-prior (Zellner, 1996) to the NN setting.
This prior allows us to obtain more expressive predictive posteriors, shown in Figure 1.
We validate our proposals, first on a 1D toy problem, and then on an image classification task, where we evaluate in and out of distribution (OOD) performance.
2. Preliminaries
We consider a supervised learning setting with observations ( x i , y i ) ni
=1 , x i ∈ R
d x and y i ∈
R d y .
We estimate y i from x i using a NN f ( x i , w ) with parameters a column vector w ∈
R d w .
To do so, we maximise an objective of the form G f ( w ) + R ( w ), where G f ( w ) =
P i ≤ n g ( y i , f ( x i , w )) is a data fit term and R ( w ) is a regulariser that encourages the entries of w to be small.
We denote by w ⋆ the output of some stochastic gradient optimisation algorithm (SGD) applied to this objective.
Within the Bayesian framework, g and R are obtained by assuming a probabilistic model for which g ( y i , ˆ y i ) is the log-likelihood log p ( y i | ˆ y i ) and R ( w ) is the log prior log π ( w ).
Thus, w ⋆ is usually understood to be a maximum a posteriori (MAP) solution or mode of the resulting posterior p ( w | ( x i , y i ) n i =1 ).
Linearised Laplace Inference for NNs The linearised Laplace method for approximating the intractable posterior is based on a 1 st order Taylor expansion of f around w ⋆
[Equation]
where J ( x ) =: ∇ w f ( x, w ) | w =
w ⋆ ∈ R d y × d w is the Jacobian of f with respect to the parameters.
The prior over v ∈ R d w is chosen to match that of the NN weights.
The posterior is then modelled as a Gaussian with mean µ = w ⋆ and covariance matrix Σ
=
( H + Λ) − −−−    w ∗ / || w ∗ || −−−    v ∗ / || v ∗ ||
NN log likelihood −−−    w ∗ / || w ∗ ||
NN log prior −−−    w ∗ / || w ∗ ||
NN log posterior where Λ = −∇ 2 w log π ( w ) | w = w ⋆ ∈ R
d
w × d w is the Hessian (or precision) of the log prior and H = X i
[Equation]
n
[Equation]
where
[Equation]
the Hessian of the log-likelihood.
We assume Λ is positive definite and let H be positive semi-definite.
Under linearised Laplace, E [ h ( x, v )]
= f
( x, w ⋆ ) and Var[ h ( x, v )]
=
∥ J ( x ) ∥ 2 Σ .
Laplace Model Evidence Approximation We adopt the common Gaussian prior
N (0 , Λ − 1 ).
The precision Λ is usually chosen to maximise the volume of a Gaussian approximation to the posterior mode containing w ⋆ .
This so-called model evidence is given by
[Equation]
where C is independent of Λ, and w ⋆ is the linearisation point obtained with an optimisation algorithm.
The norm term encourages the prior precision Λ to be small enough to accommodate w ⋆ .
The determinant ratio encourages larger precisions such that the contraction of the posterior’s variance relative to the prior’s is minimised.
3. Pathologies in Linearised Laplace with Normalised Networks
We suppose there exists a partition of the neural network parameters w into L + 1 groups w 1 , . . .
, w L +1 , e.g. layers, and that the output of the network is invariant to the scaling of groups of parameters w 1 , . . .
, w L .
That is, taking R + = { x ∈ R : x > 0 } , for all k ∈ R L + ,
[Equation]
where k · w 1: L = S Li =
1 k
i w i refers to set-element-wise multiplication in the rest of the paper.
Figure 2 illustrates this radial invariance.
We call a neural network satisfying (4) a normalised neural network .
This invariance occurs in, for example, layer norm (Ba et al., 2016), group norm (Wu and He, 2020), batch norm (Ioffe and Szegedy, 2015), or even in so-called normalisation-free methods (Brock et al., 2021).
Group w L +1 , which corresponds to the output layer of the neural network, is not invariant.
−−−    w ∗ / || w ∗ ||
−−−    v ∗ / || v ∗ ||
Linearised log likelihood −−−    w ∗ / || w ∗ ||
Linearised log prior −−−    w ∗ / || w ∗ ||
Linearised log posterior The projection is onto these two directions. The posterior mode v ⋆ is different from w ⋆ . Recall that G f ( w ) =
P i ≤ n log p ( y i | f
( x i , w )) and π ( w ) =
N ( w ; 0 , Λ − 1 ).
For f , a normalised neural network, a maximum of G f ( w ) + log π ( w ), also known as a maximum a posteriori, does not exist .
To see this, suppose w ⋆ is a MAP.
Take w ′ =
w ⋆L +1 ∪ (0 . 5 · w ⋆ 1: L ).
Then
G f ( w ′ )
=
G f ( w ⋆ ) and log π ( w ′ ) >
log π ( w ⋆ ) (5) and thus w ⋆ is not a MAP solution.
Additionally, let w ⋆l be a normalisation group of w ⋆ (as used in eq. (4)), with l ≤ L , and corresponding Jacobian entries J l ( x ) = ∇ w l f ( x, w ) | w = w ⋆ .
Then w ⋆l lies in the null space of J l ( x ).
This can be seen by taking the directional derivative
[Equation]
Similar results have been noted in the optimisation literature, where the prior term is introduced through weight decay (van Laarhoven, 2017; Hoffer et al., 2018;
Cai et al., 2019;
Li et al., 2020; Lobacheva et al., 2021).
As a consequence of eq. (6), the gradient of G f acts like a tangential force, making the norm of w larger.
The prior gradient acts like a centripetal force, keeping the norm small.
This can be seen in Figure 2.
van Laarhoven (2017) shows that SGD leads to a norm at convergence of ∥ w ⋆ ∥ 22 ∝ ( σ ( ∇ w G f ) / λ ) 0 . 5 , where σ ( ∇ w G f ) is the standard deviation of ∇ w G f
and λ is the precision of an isotropic Gaussian prior.
This can be understood as SGD returning w L +1 ∪ k · w 1: L , where k ∈ R L + depends on extraneous factors such as the learning rate or batchsize.
How does the above pertain to linearised Laplace?
Denote by G h ( v ) =
P i ≤ n log p ( y i | h
( x i , v ))
the log likelihood function of the linearised model eq. (1), and recall that π ( v ) =
N ( v ; 0 , Λ − 1 ).
The linearisation point w ⋆ is not the maximum of G h ( v ) + log π ( v ).
To see this, suppose that w ⋆ is the MAP and take w ′ to be w ⋆L +1 ∪ (0 · w ⋆ 1: L ).
We have
[Equation]
thus
G
[Equation]
and log
[Equation]
and w ⋆ is not a MAP solution for the linearised model.
The Laplace method relies on w ⋆ being a posterior mode such that the first order term in a quadratic expansion of the posterior vanishes.
As a consequence of eqs.
(5) and (7), the Laplace evidence from eq.
(3) does not approximate the volume of a posterior mode for neither the NN nor linearised NN , making it unsuitable for model selection.
In practice, maximising eq. (3) leads to arbitrary results due to the dependence on the norm of the SGD solution, as we will see in Section 6.
Why do Daxberger et al.
(2021b) see such strong empirical performance with linearised Laplace in normalised NN architectures despite the above?
The key insight is that the predictive variance Var[ h ( x, v )]
=
∥ J ( x ) ∥ 2Σ , does not depend directly on the norm of w ⋆ . Because Daxberger et al.
(2021b)
employ cross-validation, as opposed to the model evidence (MLL), to select Λ, they do not encounter the issues described above.
In summary, SGD provides us with a solution w ⋆ of arbitrary norm which may be a mode of the likelihood but is not a mode of the posterior (MAP solution) for either the normalised NN or its corresponding (tangent) linear model.
Given a suitable prior precision, this does not create issues when estimating uncertainty because the posterior predictive variance is invariant to the norm of the linearisation point w ⋆ , only depending on its angle.
However, normalisation breaks the Laplace MLL objective, which relies on the norm of w ⋆ as a measure of model complexity.
4. Finding a MAP to the Lost Linearised Evidence
We proceed to obtain the evidence for the linearised model corresponding to a normalised network.
Its posterior landscape is shown in Figure 3.
First, we note that for a NN normalised as in eq. (4), and with a fully connected output layer, the first order expansion matches a simple basis function linear model
[Equation]
The final equality is due to the Jacobian entries corresponding to the last layer weights matching the last layer activations: f ( x, w ⋆ ) =
J L +1 ( x ) w ⋆L +1 .
Referring to the MAP solution of G h ( w ) + π ( w ) as v ⋆ , we write the linear model’s Laplace MLL objective as log Z Y i ≤ n
[Equation]
By eq.
(7) we know that v ⋆ ̸ = w ⋆ .
From eq.
(8) we see that v ⋆ corresponds to the MAP slope of a regression hyperplane on the Jacobian basis.
Intuitively, the prior density π ( v ⋆ ) can be interpreted as a measure of model complexity, unlike π ( w ⋆ ).
The large size of our NN’s weight space or use of non-linear linking functions often precludes finding v ⋆ in closed form.
Instead, we propose algorithm 1 for evaluating the gradient ∇ v log p ( y | h ( x, v )) without explicitly computing Jacobians.
We derive it in Appendix B.
Algorithm 1: evaluation for linearised model
Inputs: Neural network f , Observation x , Linearisation point w ⋆ , Weights to optimise v , Likelihood function p ( y |· ), Machine precision ε
// Set FD stepsize (Andrei, 2009)
[-≈ f ( x,w ⋆ + δv ) − f ( x,w ⋆ − δv ) 
-]
Two sided FD approximation to Jvp
// Evaluate gradient of loss at ˆ y =
J ( x ) v
// Project gradient with backward mode AD
Output: ∇ v log p ( y | h
( x, v ))
=
g v
[Figure]
5. The Neural Network g-Prior
Regardless of normalisation layers, we observe different scales in the components of the Jacobian basis, e.g. see Figure 4.
A single prior precision will be too restrictive for weights corresponding to small basis components, preventing these from influencing model predictions and error-bars while leaving weights corresponding to large components underspecified.
This issue is well studied in Bayesian linear models (Minka, 2000), with a standard solution being the use of Zellner’s g-prior (Zellner, 1996).
This is a Gaussian centred atwith precision the scaled Fisher information matrix π g ( w ) =
N ( w ; 0 , g · I − 1 ).
The scaling factor g gives name to the prior.
For NNs the Fisher is often singular (Watanabe, 2007), making it ill-suited as a prior.
We bypass this issue by diagonalising the prior
[Equation]
[Equation]
[-LL   90 120 
-]
with diag( I ( w ′ ))
= 1 / n P i<n E p
( y | f
( x i ,w ′ ))
[( ∇ w log p ( y | f
( x i , w )) 2 | w = w ′ ] I .
This matrix will vary depending on the point w ′ around which it is computed, resulting in an un-normalised density.
The null space of eq. (10) only corresponds to weights which have no effect on the output (e.g. those corresponding to dead ReLUs).
We drop these from our probabilistic model.
The Laplace approximation reverts eq. (10) to N ( w ; 0 , g · diag ( I ( w ⋆ )) − 1 ) with I ( w ⋆ ) matching the linearised NN’s Fisher.
Applying the proposed prior to our linearised model can be interpreted as downscaling our Jacobian features by their second moment while keeping our prior as an isotropic Gaussian of variance g. Details are in Appendix D.
6. Experiments
Our procedure is: 1) train NNs with standard methods, 2) place either an isotropic Gaussian
N (0 , λI ) or NN g-prior N (0 , g · diag( I ( w ⋆ )) − 1 ) over the parameters, 3) optimise the hyperparameters λ
and g using either eq.
(3) (with NN weights w ⋆ ) or eq.
(9) (with linear weights v ⋆ ), and 4) predict using linearised Laplace.
Implementation details are in Appendix E.
We first train a 2 hidden layer, 50 hidden unit, MLP with residual connections and layer norm on the “matern” 1d regression dataset from Antoran et al.
Figure 2 shows that the likelihood is radially invariant while the prior is rotationally invariant.
Once the network is linearised, the likelihood invariance only remains in the direction of w ⋆ , as shown in Figure 3.
Here we see that the norm of w ⋆ is much larger than that of v ⋆ .
Consequently, [Figure] optimising the MLL eq.
(3) returns too small a precision and larger in Figure model MLL eq.
(9) produces much more sensible results qualitatively.
the the informed and
Next, we train a ∼ 46k parameter batch norm Lenet-style CNN on MNIST.
Figureshows a much stronger test log-likelihood when hyperparameters are optimised with our proposed linearised MLL eq.
(9) than when using eq.
(3), for both isotropic and g priors.
Given that in all settings, the predictive mean is the same across methods, we can attribute the increase in LL to better calibrated error-bars.
Interestingly, our proposed model selection criteria eq.
(9) tends to choose the smallest precision that maximises the test LL.
This maximises uncertainty while not sacrificing in-distribution performance.
We evaluate OOD uncertainty estimation on rotated digits.
The hyperparameters found with the NN weights w ⋆ perform well OOD at the cost of overestimating uncertainty in-distribution.
Our proposed criteria combined with the isotropic Gaussian prior performs well in-distribution but is overconfident OOD.
The more expressive g-prior, when combined with eq. (9), retains in-distribution performance while expressing increased uncertainty OOD. 7. Conclusion
We have highlighted and provided solutions to two pitfalls of the naive application of linearised Laplace to modern NNs.
First, normalisation layers preclude finding an optima of the loss for a NN, or its tangent linear model, by optimising the loss function of the NN.
This invalidates the assumption that the point at which we linearise our model is stationary.
However, every linearisation point implies an associated basis function linear model. As we use this model to provide errorbars, we propose to also choose hyperparameters using this model’s evidence. This can be done by solving a convex optimisation problem, much simpler than NN optimisation. Second, the scales of basis functions corresponding to the Jacobian entries for different weights can be very heterogeneous.
We propose the extension of the g-prior to the NN setting.
This empirical prior assigns a variance to each weight inversely proportional to the scale of its associated feature.
Acknowledgments
The authors would like to thank Marine Schimel for helpful discussions.
JA acknowledges support from Microsoft Research, through its PhD Scholarship Programme and from the
EPSRC.
ED acknowledges funding from the EPSRC and Qualcomm.
JUA acknowledges funding from the EPSRC and the Michael E. Fisher Studentship in Machine Learning.
This work has been performed using resources provided by the Cambridge Tier-2 system operated by the University of Cambridge Research Computing Service (http://www.hpc.cam.ac.uk) funded by EPSRC Tier-2 capital grant EP/T022159/1.
Appendix A. On Invariances Introduced by Normalisation Layers
We relate eq. (4) to common normalisation layers.
For some input z ∈ R d z these apply the function
[Equation]
where the expectation and variance are estimated empirically and γ ∈ R , β ∈ R .
We have invariance to scaling the input by some constant κ ∈ R
[Equation]
For batch norm , the input to the above function z , corresponds to the different values of a neural activation across a mini-batch of data.
As a result, normalisation groups correspond to the fan-in weights of each node.
For layer norm , z corresponds to a layer’s activations for a single input.
Here, the normalisation group is the complete set of layer weights.
Group norm applies the same operation as layer norm but across a subset of a convolutional layer’s channels.
Here, the normalisation group is the subset of the previous layer’s weights which acts on any of the output channels that fall into a specific group.
Appendix B. Derivation and Analysis of Algorithm
We denote NN Jacobians as J ( x ) =: ∇ w f ( x, w ) | w = w ⋆ ∈ R d y × d w and a linear model weight vector v ∈ R d w .
d y is the output size.
Consider the linear model J ( x ) v where we wish to optimise v according to the objective log p ( y | J ( x ) v ).
Expanding the expression of the gradient using the chain rule
[Equation]
we see that it is equal to the gradient of the loss projected onto the weights through the
Jacobian.
We first need to evaluate our basis function linear model ˆ y =
J ( x ) v .
This
Jacobian vector product can be computed with forward mode automatic differentiation or finite differences (FD).
We rely on finite differences because it is slightly cheaper to compute.
Specifically, we employ the method of Andrei (2009) to select the optimal step size.
We empirically evaluate this approach in Figure 6, finding it to return the nearoptimal step-size for a wide range of queries.
We then evaluate the loss gradient at the linear model output, g y = ∇ ˆ y p
( y | ˆ y ) | ˆ y =
J ( x ) v .
This can often be done in closed form.
Finally, we project onto the weights using backward mode automatic differentiation g w = g Ty J ( x ) = ∇ w ( g Ty f ( x, w ))
| w = w ⋆ .
For the small MLP used in the toy experiments from Section 6, it is tractable to compute the Jacobians for the full training set.
In this setting, we can evaluate the fidelity of the gradient update computed with algorithm 1 relative to exact gradient descent (GD).
In
Figure 7, we show that the proposed FD based method tracks the GD trajectory very closely.
The gradient bias, computed in terms of rmse across all weights, stays below 2 e −during the whole optimisation.
The bias in the weights accumulates throughout training, reaching a maximum of ∼ 0 .
015 at convergence.
[Figure]
[Figure] In Figure [Figure] we compare the fits and weight histograms obtained when optimising the linear model corresponding to the MLP from Section 6 with algorithm 1 and with the closed form solution for Gaussian linear regression.
Both fits agree seemingly perfectly in the range of the data, with some slight disagreement in the extrapolation regime.
The weight histograms overlap almost perfectly in the body, with some outliers being placed in different locations.
This suggests that the MLL values obtained with eq. (9) when finding v ⋆ with algorithm 1 will be a good approximation to the true linearised model MLL
[Figure]
Appendix C. Additional Plots
In Figure 9 we compare the fits obtained on our toy problem when employing the NN function f ( x, w ⋆ ) and the linearised model h ( x, v ⋆ ).
In the latter case we compare both an isotropic Gaussian prior and the a g-prior.
All three methods return very similar mean predictions, with the most notable difference being the g-prior’s additional non-smoothness in the out-of-distribution regime.
This is due to the g-prior allowing for more basis functions to contribute the the mean prediction.
The histogram plot shows that w ⋆ presents a both [Figure] [Figure] wider and heavier tailed distribution than v ⋆ .
This explains why choosing hyperparameters with eq. (3) leads to excessively wide error bars in Figure 1.
The weights obtained when using the g-prior present a wider distribution than w ⋆ but show less heavy tails.
Figure 10 shows the weight histogram for the CNN used in Section 6.
Its architecture is described in
Appendix E.
The NN weights w ⋆ are significantly larger than those of the linearised model v ⋆ .
As a result, the MLL that uses w ⋆ favours smaller prior precisions which result in strong underconfidence in-distribution, as found in Section 6.
− 1.5 − 1.0 − 0.
density 10 −10 −10 −101010w
e i g h t s w ?
v ?
Appendix D. Feature Normalisation and the Diagonal g-Prior
In this appendix, we provide some intuition for the effects of the diagonal g-prior on the linearised NN model.
We stack the features ( J ( x i ))
i ≤ n , into the tensor J ∈ R n × d y ×
d w .
We will suggestively refer to the tensor that has the first and third indexing dimensions flipped as J T ∈ R d w × d y × n for notational convenience.
The log-likelihood Hessian, stacked across observations is H ∈ R n × d y × d y .
For exponential family distributions, the Fisher I ∈ R d w × d w matches the GGN matrix (Martens, 2014)
[Equation]
(11) and thus the NN g-prior precision is
[Equation]
A ∈ R d w × d w is a diagonal matrix which will act as a feature normaliser J s =
J A − 1 .
The entries of A are
[Equation]
(13) where H regulates linear interactions across output dimensions and observations.
D.1. Parameter Optima, Linearised Laplace Predictive Distribution and Laplace Marginal Likelihood induced by Feature Normalisation and Diagonal g-Prior
Consider a generalised linear model h ( x, w ), which we optimise with the loss function G (ˆ y ) =
P i ≤ n log p ( y i | ˆ y ).
We consider 2 scenarios.
In the first we choose the regulariser based on the diagonal g-prior log π g ( v ) ∝ − 0 .
5 g − 1 · v T A T Av .
In the second we choose an isotropic regulariser log π ( v ) ∝ g − 1 · v T v
but we scale our features as J s =
J A − 1 .
We refer to the former as the g-prior model and to the latter as the scaled model.
Parameter Optima
An optima of the parameters v ⋆ satisfies
[Equation]
We first apply the chain rule to the loss gradient ∇ w G ( w ) | w = w ⋆ = ∇ ˆ y G (ˆ y ) | ˆ y =
w ⋆ J J .
In the g-prior setting, this is
[Equation]
[Equation] and we have denoted the stationary point v ⋆g .
For the scaled model, the stationary point is
[Equation]
with stationary point v ⋆ s .
Combining both eqs.
(15) and (16) we get that the scaled model weights are just the scaling matrix applied to the g-prior weights A − 1 v ⋆ s = v ⋆g .
We now show this in the special case of the Gaussian likelihood case, with precision H , v ⋆g = n X i
[Equation]
Linearised Laplace Predictive
The predictive mean for both the g-prior model and scaled model is the same:
[Equation]
The Laplace predictive covariance is also the same:
[Equation]
Linearised Laplace MLL
This quantity also matches for the g-prior model and scaled model:
[Equation]
We have dropped the data fit terms and other constants as these match for the g-prior and scaled model since both models make the same mean predictions eq.
[Equation]
D.2. What type of Feature Normalisation does the Diagonal g-Prior Induce?
In this subsection we discuss the similarity between the scaled model, induced by use of the g-prior and a generalised linear model in which the inputs have been scaled such that all inputs share some summary statistic (e.g. 0 mean, unit variance, etc).
We make the assumption that our data is iid and thus the likelihood function is factorised across datapoints.
Let’s first consider the case where our log-likelihood is homoscedastic and factorised across output dimensions H i = κ · I ∀ i ≤ n , where κ ∈ R + is the likelihood precision.
Here we have that the scaling matrix matches the empirically computed second moment of each basis function in our expansion
[Equation]
summed across output dimensions and scaled by the noise precision.
In fact, the constant κ can be absorbed into g .
For one output dimension d y = 1, this leaves us with a scaling procedure that matches commonly used standard deviation normalisation for the case of 0mean features.
Whether our Jacobian features are 0-mean will depend on if the linearisation point w ⋆ is stationary.
In the multi-response setting d y > 1, the standardisation for the feature corresponding to each weight is proportional to the sum of the second moments of the gradients of each output dimension with respect to the weight.
For a heteroscedastic likelihood function factorised across output dimensions H i
= κ i · I ∀
i [Equation] n , where κ ∈ R n + we have
[Equation] [Equation]
Now different datapoints contribute to the scaling factor proportionally to their noise precision.
Finally we consider the case where our likelihood function is heteroscedastic and does not factorise across output dimensions, like is the case for the softmax-categorical.
Here H i is full rank.
The entries of our scaling matrix are given by eq.
Appendix E. Image Classification Experimental Setup
For the image classification experiments, we employ a CNN based on the LeNet architecture with a few variations found in more modern neural networks.
The architecture contains 3 convolutional blocks, followed by global average pooling in the spatial dimensions, a flatten operation, and finally a fully-connected layer.
The convolutional blocks consist of a convolution layer, a ReLU activation, and a batch norm layer, in that order.
Instead of using max pooling layers, as in the original LeNet variants, we use convolutions with a stride of 2.
The first convolution is 5 × 5, while the next two are 3 × 3.
Table 1 shows the sizes of the filters and number of parameters.
These where chosen to create a model as large as possible while keeping full-covariance Laplace inference tractable on one A100 GPU.
The NN weights w ⋆ are learnt using SGD, with an initial learning rate of 0.1, momentum of 0.9, and weight decay of 1 × 10 − 4 .
We trained for 90 epochs, using multi-step LR scheduler with a decay rate of 0.1 applied at epochs 40 and 70.
[Table]
The linear weights v ⋆ are learnt using SGD but with the gradient calculated via algorithm 1.
We use a learning rate of 1 × 10 − 4 and train for 100 epochs.
We set the weight decay value to 1 × 10 − 4 when using isotropic priors and to 1 × 10 − 1 when scaling features according to the g-prior.
The latter choice is made due to the model with scaled features having a larger effective dimensionality and thus needing stronger regularisation.
We do not use momentum or learning rate decay.