Accelerated Linearized Laplace Approximation for Bayesian Deep Learning
Authors: Zhijie Deng
Feng Zhou
Jun Zhu
Venue: NeurIPS 2022 Accept (2022)
Keywords: Bayesian deep learning
Laplace approximation
kernel approximation

Abstract: Laplace approximation (LA) and its linearized variant (LLA) enable effortless adaptation of pretrained deep neural networks to Bayesian neural networks. The generalized Gauss-Newton (GGN) approximation is typically introduced to improve their tractability. However, LA and LLA are still confronted with non-trivial inefficiency issues and should rely on Kronecker-factored, diagonal, or even last-layer approximate GGN matrices in practical use. These approximations are likely to harm the fidelity of learning outcomes. To tackle this issue, inspired by the connections between LLA and neural target kernels (NTKs), we develop a Nystrom approximation to NTKs to accelerate LLA. Our method benefits from the capability of popular deep learning libraries for forward mode automatic differentiation, and enjoys reassuring theoretical guarantees. Extensive studies reflect the merits of the proposed method in aspects of both scalability and performance. Our method can even scale up to architectures like vision transformers. We also offer valuable ablation studies to diagnose our method. Code is available at https://github.com/thudzj/ELLA.

Accelerated Linearized Laplace Approximation for Bayesian Deep Learning
Zhijie Deng 2, Feng Zhou 23, Jun 24
Qing Yuan Research Institute, Shanghai Jiao Tong University 2 Dept. of Comp. Sci. & Tech.. BNRist Center, THU-Bosch Joint ML Center, Tsinghua University
3 Center for Applied Statistics, School of Statistics, Renmin University of China Pazhou Laboratory (Huangpu), Guangzhou, China edu. feng. zhou@ruc. edu. en,
edu. cn
Abstract
Laplace approximation (LA) and its linearized variant (LLA) enable effortless adaptation of pretrained deep neural networks to Bayesian neural networks. The generalized Glaus--Netoon (GGN) approximation is typically introduced im- prove their tractability. However, LA and LLA are still confronted with non-trivial inefficiency issues and should rely on Kronecker-factoree layer approximate GGN matrices in practical use. These approximations are likely harm the fidelity of learning outcomes. To tackle this issue, inspired by the connections between LLA and neural tangent kernels (NTKs), we develop Nys- tröm approximation NTKs to accelerate LLA. Our method benefits from the capability of popular deep learning libraries for forward mode automatic differ- entiation, and enjoys reassuring theoretical guarantees. Extensive studies reflect the merits of the proposed method in aspects of both scalability and performance. Our method can even scale up to architectures like vision transformers. We also offer valuable ablation studies to diagnose our method. Code is available at https: //gi thub. wwt.cto////o//. nom/thudzj /ELLA.
diagonal, even last-
Introduction
Deep neural networks (DNNs) excel at modeling deterministic relationships and have become de facto solutions for diverse pattern recognition problems [18, 54]. However, DNNs fall short in reasoning about model uncertainty [1] and suffer from poor calibration [17]. These issues are intolerable in risk-sensitive scenarios like self-driving [27), healthcare [34], finance [25], etc.
Bayesian Neural Networks (BNNs) have emerged as effective prescriptions to these pathologies [37, 21,44, 16). They usually proceed by estimating the posteriors over high-dimensional NN parameters. Due to some intractable integrals, diverse approximate inference methods have been applied to learning BNNs, spanning variational inference (VI) [1, 36], Markov chain Monte Carlo (MCMC) [55, 4], Laplace approximation (LA) [38, 50, 30], etc.
LA has recently gained unprecedented attention because its post-hoc nature nicely suits with the pretraining finerrunnn fashion in deep learning (DL). LA approximates the posterior with Gaussian around its maximum, whose mean and covariance are the maximum posteriori (MAP) and the inversion of the Hessian respectively. It is common practice to approximate the Hessian with the generalized Gauss-Newton (GGN) matrix [40] make the whole workflow more tractable.
Linearized LA (LLA) [14, 22] applies LA to the first-order approximation of the NN of concern. Immer et al. [22] argue that LLA is more sensible than LA the presence of GGN approximation:
"Coresponding author. Majority of work completed while Zhijie Deng was at Tsinghua University.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

11A--DE
I-D regression on y sin 2x e, (0, 0.2). Red dots, central bluc curves, and shaded regions refer to the training data, mcan predictions, and uncertainty respectively. The model is pretrained multilayer perceptron (MLP) with 3 hidden layers. As shown, the predictive uncertainty of ELLA is on par with or better than the competitors such as LLA with KFAC approximation (LLA-KFAC), LLA with diagonal approximation (LLA-Diag), and last-layer LLA (LLA").
Figure I:
LLA can perform on par with or better than popular alternatives on various uncertainty quantification (UQ) tasks [14, 6, 7]. The Laplace library [6] further substantially advances LLA's applicability, making it a simple and competing baseline for Bayesian DL.
The practical adoption of LA and LLA actually entails further approximations top of GGN. E.g., when using Laplace to process a pretrained ResNet [18], practitioners are recommended to resort to Kronecker-factorde (KFAC) [40] or diagonal approximation of the full GGN matrix for tractability, An orthogonal tactic to apply LA/LLA to only NNs' last layer [30]. Yet, the approximation errors in these cases hardly be identified, significantly undermining the fidelity of the learning outcomes.
This paper aims at scaling LLA up to make probabilistic predictions in more assurable way. We first revisit the inherent connections between Neural Tangent Kernels (NTKs) [24] and LLA [28, 22], and find that, if we can approximate the NTKs with the inner product of some low-dimensional vector representations of the data, L.A can be considerably accelerated. Given this finding, we propose to adapt the Nyström method to approximate the NTKs of multi-output NNs, and advocate leveraging forward mode automatic differentiation (fwAD) to efficiently compute the involved Jacobian- -vector products (JVPs). The resultant accElerated LLA (ELLA) preserves the principal structures of vanilla LLA yet without explicitly computing/storing the costly GGN/Jacobian matrices for the training data. What's more, we theoretically analyze the approximation error between the predictive of ELLA and that of vanilla LLA, and find that it deceases rapidly as the Nyström approximation becomes accurate.
We perform extensive studies to show that ELLA can be a low-cost and effective baseline for Bayesian DL. We first describe how to specify the hyperparameters of ELLA, and use an illustrative regression task to demonstrate the effectiveness of ELLA (see Figure 1). We then experiment on standard image classification benchmarks to exhibit the superiority of ELL.A over competing baselines in aspects of both performance and scalability. We further show that ELLA can even scale up to modern architectures like vision transformers (ViTs) [11].
Background
{(x; v)) where æ; X' and y: RC (e.g.. regression) {0, 1 (c.g., classification) refer to observations and targets respectively. The advance in machine learning suggests using an NN ge(-) X - RC with parameters 0 H." for data fitting. Despite well-performing, regularly trained NNs only capture the most likely interpretation for the data, thus miss the ability to reason about uncertainty and are prone to overfitting and overconfidence.
Consider a learning problem
(X, Y)
BNNs [37, 21, 44] characterize model uncertainty by probabilistic principle and can holistically rep- resent all likely interpretations. Typically, BNNs impose a prior p(0) on NN parameters and chase the Bayesian posterior D) p(D ())(()))(()))))))) ()((p)))))))))pp())))))))))))( where p(D)0)
P(y. (x., 0) II, p(y, |ge (x:)). Analytical estimation is usually intractable due to NNs' high nonlinearity. Thereby, BNN methods usually find a surrogate of the true posterior (0) p(O (D) via approximate inference methods like variational inference (VI) [1. 20, 36, 62. 29], Laplace approximation (LA) [38, 50], Markov chain based variational inference (POVT) [35], etc.
Monte Carlo (MCMC) [55, 4, 63]. particle-opimmzztte
BNNs predict for new data x. by posterior predictive pyyl.., D)
Ep p)P(3) ax,, 0) R
Enco x. 0) $ E p(y go. (xx)) where O, q(0) are i.id. Monte Carlo (MC) samples.

---

2.1 Laplace Approximation and Its Linearized Variant
Typically, LA builds Gaussian approximate posterior in the form of q(0) = O, Σ), where @ de- notes the MAP solution, i.e., 0 = arg m&xo log p( log and Σ is the inversion of the Hes- sian of the negative log posterior w.r.t. parameters, i.e., Σ - (log p(D)0) + log (())) lo) Without loss of generality, p(0) = \i Ip)', so - V2 log p(p)) equals to Ip/oğ-
base the following discussion the isotropic Gaussian prior
Due to the intractability of the Hessian for NNs with massive parameters, it is common practice to use the symmetric positive and semi-definite (SPSD) GGN matrix workaround, i.e.,
-'= = Σ Jo(x.) A(a., (x.)) 1p/00:
(I)
where 10(x) & Vego(()) 0=0 and A(x, y) - V
og go g=
When concatenating (Jg( E RCXPjNN {A(x;, y:) E RRCxC Ji=1 as a block-diagonal matrix Ax,Y E IRNCXNC we have:
big matrix J J8, RNCXP and organizing
Σ [Jxx T Ax, + Lr/od]
Yet, LA suffers from underfitting [33]. This is probably because GGN approximation implicitly turns
but ge ge(x) (x) is still leveraged make prediction. Le.. there exists shift between posterior inference and prediction [22]. To mitigate this issue, a proposal is to predict with (zz), giving rise to linearized LA (LLA) [14, 28, 22].
the original model (x) into a linear one (xx)
LA and LLA nicely fit the pretraininngiiietuniing fashion in DL they can be post-hoc applied to pretrained models where only the GGN matrices require to be estimated and further inverted. By LA and LLA, practitioners can adapt off-the- shel- high performing DNNs to BNNs easily.
LLA has revealed strong results on diverse UQ problems [14, 6, 7]. The Laplace library [6] further advances I.AA's applicability,, and evidences L.AA is competitive to popular alternatives [63, 32, 39].
Scalability issue The GGN matrix of size Px P is still unamenable in modern DL scenarios, so further approximations sparsifying it are always introduced. The diagonal and KFAC [40] approximations are commonly adopted ones [50, 62), where only diagonal block-diagonal structure the original GNN matrix is preserved. An orthogonal tactic is to concern only subspace of the high-dimensional parameter space (e.g., the parameter space of the last layer [30]) to reduce the scale of the GGN matrix. However, these strategies sacrifice the fidelity of the learning outcomes as the approximation errors these cases can hardly be theoretically measured. To this end, we develop accElerated Linearized Laplace Approximation (ELJA) to push the limit of LLA in a more assurable way.
3 Methodology
In this section, we first revisit the relation of LLA Gaussian processes (GPs) and Neural Tangent Kernels (NTKs) [24]. After that, we reveal how to accelerate LLA by kernel approximation. Based on these findings, we develop an efficient implementation of ELLA using the Nyström method [58].
3.1 The Gaussian Process View of LLA
Integrating (0) = (0 ; Ó, E) with the linear model g (a) actually gives rise to a function-space approximate posterior [14, 28, 22] in the form of q(f)) = GP gj (x), (aaa x²)) with
MLLA x') e ()()) ('')
(3)
By Woodbury matrix identity [60], we have:
[Ja, Ax,YJax Ir/oo]] 0ā Ip Ja, (AxIx//0 Jax T- TJJaa))
(4)
It follows that
(x, oe oz (a) (T, x) XXTK(XX "))xXXNX((Xxx')) (5)
'I refers to the identity matrix of size P
3

---

where KYTK(Z, x') _ [ denotes the neural tangent kernel (NTK) [24] corresponding to gg(x)) Note that NNTK (z, x²) is a matrix-valued kernel, with values in the space of C' x C matrices.
The main challenge then turns into the computation and inversion of the gram matrix ANTK X X) of size NC NC. When either N or C' is large, the estimation of KLLA still suffers from inefficiency issue. To address this, existing work [22] assumes independence among the C' output dimensions cast q(f) into C' independent GPs following [51], and randomly subsample M < N data points to form a cheap substitute for the original gram matrix. Despite effective in some cases, these approximations are heuristic, lacking a clear theoretical foundation.
3.2 Scale Up LLA by Kernel Approximation
We show that if we can approximate KNTK (ar, æ²) with the inner product of some explicit C x K- dimensional representations of the data, i.c. KNTK (x, x²) ((x))(((') with p : X CxC the scalability of the whole workflow can be unleashed.
Concretely, letting x R (RNC be the concatenation of (1p(x)) ROCK Ii have (detailed
derivation in Appendix A.2; see also [9])
(4(() (<)))(((()² p())pp Ax.y xxxppx (xxp(x))))
=x( Σ ((x;)) A(ai, v))p(a)))))))))))))))))))))))( Ix/00] (0xx') (o,
(6)
The matrix G I k possesses a similar formula to the Σ in the original KLLA in Equation (3), yet much smaller. Once having ye, it is only required perform one forward pass of g and io for each training data to estimate G. When K is reasonably small, e.g., < 100, is cheap to invert G.
3.3 Approximate NTKs via the Nyström Method
Typical kernel approximation means include Nyström method [45, 58] and random Seatures-based methods [48, 49, 61, 43, 15, 9]. Given that the latter routinely relies on relatively large number of random features to gain a faithful approximation (which means a large K), we suggest leveraging the Nyström method, which captures only several principal components of the kernel, to build (p.
To comfortably apply the Nyström method, we first rewrite ( x²) as a scalar-valued kernel KNTK ((xr, i), i") Jg(1: j)Ja(a', ,i)" where Jg( A Va [=)( lo l a X x [C] R1X2 computes the gradient of i-th output w.r.t. parameters.
By Mercer's theorem [42].
ONTK ((x, 1). ()), = Σ (u,
where 2/ E L' x [C).g)² refer to the eigenfunctions of the NTK w.r.t some probability measure q with μ4 2 0 as the associated eigenvalues. Based on this, the Nyström method discovers the top-K eigenvalues as well as the corresponding eigenfunctions for kernel approximation.
By definition, the eigenfunctions can represent the spectral information of the kernel:
[ FNIK ((x, i), (a',"")) the (a', (')q(a', (') = pa i), Vk 2 1,
(8)
while being orthonormal under q: (ax (ar, 1) ()g(2, a) 1lk V1, Vk. 1
In our case, q(x,, i) can be factorized as the product of the data distribution and uniform distribution over {1,..., C), which can be trivially sampled from. The Nyström method draws M (M 2 K) i.i.d.
samples X = {(01 i..,.. (eMs&M)) from & to approximate the integration in Equation (8):
1
Σ
NNTK ((x, f), (æm, km)) t/x (Xm, fm) (a f), Vk [K].
(10)
M
m=
We use superscript to index a multi-dimmnnional tensor and [C] represents set integers from to C. 'L2( 4) denotes the space of equare-integrable functions w.r.t. g.

---

Applying this equation to these samples gives rise to:
M
M KNIK ((wm' kmn), (20m, ) the (2mm im) paaete ( Wk [K], [M],
Σ
(11)
then arrive at
Kip peape, [K];
(12)
M
where K J0x 0,x' with J&. RAIXP as the concatenation of {Ja (zm: im) RIIPP
mm=
and 4/x- = ( thx (ar),)))))'" This implies scaled eigendecomposition problem of K.
We compute the top-K eigenvalues ...... Ax of matrix K and record the corresponding orthonormal eigenvectors ",,..., uk. Given the constraint in Equation (9), it is easy to that:
Ak and 4/x (Zm, ss v VM (mm) m E [M]. M
A R
(13)
Combining Equation (13) and (10) yields the Nyström approximation of the top-K eigenfunctions:
VM
M
(m)) KNTK ((x, i), (Ims im)) = VM
(x)
Σ
uk
J6 (a. i)J 6.8 uk.
(14)
A
A
m=1
Then, the mapping yp which satisfied (N, x') z ((x)) can be realized
[+(x)]((A) s'x (ze pn = Ja i)J 0. uk/ / Ave
(15)
x(x)
(g(x)) with U2 J U
(16)
Connection to sparse approximations of GPs A recent work [57] has shown that sparse variational
Gaussian processes (SVGP) [53, 3] is algebraically equivalent to the Nyström approximation because the Evidence Lower BOund (ELBO) of the former encompasses the learning objective of the latter.
That's say, ELLA implicitly deploys sparse approximation the original GP predictive of LLA.
Discussion In this work, the developed Nyström approximation to NTKs mainly used for accelerat- ing the computation of the predictive covariance of LLA, while it can also be easily applied to future work on the practical application of NTKs.
3.4 Implementation
As shown, the estimation of & on data point degenerates as K JVPs, which can be accomplished by invoking fwAD for K times:
θ
forward
(x(x) Jo (z))
x.
[K].
(17)
where the model output 9g (x) and the JVP (z) are simultaneously computed one single forward pass. Prevalent DL libraries like PyTorch [47] and Jax [2] have already been armed with the capability for fwAD. Algorithm details the procedure of building in PyTorch-like style.
With that, can trivially compute the posterior q(f) according to Equation (6), as detailed in Algorithm The procedures for estimating G and the posterior can both batched for acceleration.
Computation overhead The estimation of involves M forward and backward passes of 99 and the eigendecomposition of matrix of size M x M. After that, we only need to make (w), } persistent, which amounts to storing K more NN copies. Estimating G requires scanning the training set once and inverting a matrix of K K. which cheap. The evaluation of g(f) embodies that of sp, i.c., performing K forward passes under the scope of fwAD. This is similar to other BNNs that perform S forward passes with different MC parameter samples to estimate the posterior predictive.
4 Theoretical Analysis
In this section, we thcoretically analyze the approximation enror of KELLA (w. (x,x') to KLLA KLLA(E, (x, x').
5

---

Algorithm 1: Build the LLA posterior.
Algorithm 2: Build
I6 NN pre-trained by MAP; (X, Y): training set; C: number of classes
def b X,C,M,K); def _g(0, C, {uk (== .)) Po zezorsrC K) for in range (K):
M, M,K,00 hyper-parameters def estimate ( (o,))
G =zeros (K, K)
with fwAD, fwdD.enane enable
for y) in (X,Y):
gm:jup = (() (z)
ge. = (x)
Pa (p =
Az,y = hessian (nll
return 9r: L
G + Az
Ja. =zero (M (i))
return G+ eye eye(K))))))) (K)
for m in range(M()
def (p, G
um uniforr_ssmpplelee (X)
Ja: Po = y (a:)
in = m.m=forr__sspppleee (C)
R2,= G 4p
Paa
Jg [ax(m]=grad((p (Tm [im])
return N (gz: Rzzzz)
{as, un 1 eig (J x 8 top = K)
& (gg X, C, M, K)
for in range(K),
G =inv(estimate_ G (p, X, Y, ()))
8 x V
f pa
return partial Caturnpartiil(C.mm 9o C,,
To reduce unnecessary complexity, this section assumes using M K i.i.d. MC samples when per- forming the Nyström method. Then, KELLA can be reformulated as (details deferred to Appendix A.3)
( x') (x))) Jb.2 Jo Ax o. Jax JJ Jo Jox Jo(x))
(18)
Thereby, the gap between KELLA (z, æ') and KLLA (æ, x²) can be upper bounded by & = | - E|| where || represents the matrix 2-norm (i.e., the spectral norm). In typical cases, P > M, thus O. present an upper bound of as follows.
with high probability K = Ja 8.
Theorem (Proof in Appendix A.4). Let CA be a finite constant associated with A, and E' the error
of Nyström approximation | x (J Jox
1oxx - JJxx 0. Jo J It holds that 6.x
& S cf + +
E' has been extensively analyzed by pioneering work [12, 5, 26), and we simply adapt the results developed by Drineas and Mahoney [12] to case. We denote the maximum diagonal entry and the M + 1-th largest eigenvalue of J J J by Cx and Asst1 respectively.
Theorem (Error bound of Nyström approximation). With probability least 1 5, it holds that:
NC
1 (2 (2 log 8)
E' < AM+t
VM
Plugging this back to E, we arrive at the corollary below.
Corollary 1. With probability least -8, 6, the following bound exists:
NC
1 )) oØ. (2 + log ))
E ofea
VM
As desired, the upper bound of & decreases along with the growing of the number of MC samples M.
Experiments
We first discuss how to specify the hyperparameters of ELLA,, then expose an interesting finding. After that, we compare ELLA to competing baselines to evidence its merits in efficacy and scalability.

---

0.25
14
11.6
1124
255
0.22
ReaNc32
RecNe20
512
0.29
ResNest4
ResNe35
001
25 5u
16n
0
10000
20000
30020
40000
(a)
(b)
(c)
Figure 2: (a)-(b): The approximation errors 6Nyarö and CELLA vary w.r.t M and K (M 2 K). (c): The test NLL of ELLA on CIFAR-I0 10 varies w.r.t. the number of training data N,
5.1 Specification of Hyperparameters
Given an NN 99 pretrained by MAP, we need specify the prior variance o the number of MC samples in Nyström method M, and the number of remaining eigenpairs K before applying ELLA. We simply of Ny with y as the weight decay coefficient used for pretraining according to [10].
We perform empirical study inspect how M and K affect the quality of the Nyström approxima- tion and the ELLA GP covariance. We take 2000 MNIST images as training set X, and 256 others validation X The architecture contains 2 convolutions and a linear head. Batch normaliza- tion [23] and ReLU activation are used. The number of parameters P is 29, 034. Larger architectures or larger will render the exact evaluation of KLLA unapproachable. We quantify the approximation [Jxx 0.x (J 8. 1 J 0.x Ja,x [X.al ExEx., ||MELLA (x. x) - ALLA (ar, m)]l/ (|ll/|| (ar, x)]l- We vary M from 4 to 2000 and K from to M, and report the approximation errors in Figure 2 (a)-(b). We notice that 1) the larger K the better; 2) when fixing K, ENyström and ÉELLA decay as M grows; 3) Nyssrrm decays more rapidly than EFILA, echoing Theorem Given that ELLA needs to store K vectors of size P. small K is desired for efficiency. K E [16, 32] seems to be reasonable choice given Figure Besides, Appendix C. includes direct study on how the test NLL of ELLA varies w.r.t. M and K on CIFAR- 10 [31] benchmark. Given these results, set M 2000 and K 20 in the following experiments unless otherwise stated.
error of the Nyström method by CNysnti
b.x I/1ll/// x x
and that from KELLA to KLLA by FELLA 1
ELLA (or LLA) finds a GP posterior, so predicts with a variant of the aforementioned pos- Epprpip(yl Ey~Niss (ee (... ..))))yyylll ()((y))) ff.. In classification tasks, we use 512 MC samples to approximate the last expectation and it is cheap.
terior predictive, formulated as p(y],x,
5.2 The Overfitting Issue of LLA
Reinspecting Equation (1) and (6), we see, with more training data involved, the covariance in LA, LLA, and ELLA shrinks and the uncertainty dissipates. However, under ubiquitous model misspecification [41], should the uncertainty vanish that fast? We perform study with ResNets [18] on CIFAR-10 10 to seek an answer. Concretely, we randomly subsample N data from the training set of CIFAR-10, and fit ELLA them. We depict the test negative likelioel (NLL), accuracy, and expected calibration error (ECE) [17] the deployed ELLA in Figure (c) and Appendix C.2. We also provide the corresponding results of LLA* in Appendix C.3. The V-shape NLL curves across settings reflect the overfitring issue of ELLA and LLA" (or more generally, LLA). Figure also shows that tuning the prior precision w.r.t. marginal likelihood can alleviate the overfitting of LLA* to some extent, which may be the reason why such an issue has not been reported by previous works. But also of note that tuning the prior precision cannot fully eliminate overfitting.
To address the overfitting issue, we advocate performing early stop when fitting ELLA/LLA on big data. Specifically, when iterating over the training data to estimate G (see Equation (6)),
'We experiment on LLA' here due to its higher efficiency than LLA KFAC and LLA-Diag. LLA" is the default option in the Laplace library [6].
7

---

Table Comparison on test accuracy (%) 1, NLL 4, and ECE on CIFAR-10. report average results over 5 random runs. As the accuracy values of most methods are close, we do not highlight the best.
ResNet-20
ResNet-32
ResNet-44
ResNet-56
Method
Acc.
NLL
ECE
Acc.
NLL
BCE
Acc.
NLL
ECE
Acc.
NLL
ECE
ELLA
92.5
0.233
0.009
93.5
0.215
0.008
93.9
0.204
0.007
94.4
0.187
0.007
MAP
92.6
0.282
0.039
93.5
0.292
0.041
94.0
0.275
0.039
94.4
0.252
0.037
MFVI-BF
92.7
0.231
0.016
93.5
0.222
0.020
93.9
0.206
0.018
94,4
0.188
0.016
LLA"
92.6
0.269
0.034
93.5
0.259
0.033
94.0
0.237
0.028
94.4
0.213
0.022
LLA*-KFAC
92.6
0.271
0.035
0.260
0.033
94.0
0.232
0.028
94.4
0.202
0.024
LLA-Diag
92.2
0.728
0.404
92.7
0.755
0.430
92.8
0.778
0.445
92.9
0.843
0.480
LLA-KFAC
92.0
0.852
0.467
91.8
1.027
0.547
91.4
1.091
0.566
89.8
1.174
0.579
ELLA
F.L.A.
MAM
NEVAHP
Skot
Figure NLL (Lefft) and ECE (Right) on CIFAR- corruptions for models trained with ResNet-56 architecture. Each hox corresponds to summary the results across types skew.
continuously record the NLL of the current ELLA posterior some validation data, and stop when there is trend of overfitting. If we cannot access a validation set easily, we can apply strong data augmentation some training data form a substitute of the validation set. Compared tuning the prior, early stopping also helps reduce the time cost of processing big data (e.g., ImageNet [8]).
5.3 Illustrative Regression
We build a regression problem with N 16 samples from y sin 2x +,, (0,00.2) as shown in Figure 1. The model is an MLP with 3 hidden layers and tanh activations, and we pretrain it by MAP for 1000 iterations. For ELLA, we M 16 and K 5 for efficiency. Unless stated otherwise, we use the interfaces of Laplace [6] to implement LLA, ALA-KFAC, LLA-Diag and LLA" The hyperparameters of the competitors are equivalent those ELLA except for some dedicated ones like M and K. It is clear that ELLA delivers a closer approximation to LLA than LLA-Diag and LLA* We further quantify the quality the predictive distributions produced by these approximate LLA methods using certain metrics. Considering that in this case, the predictive distribution for one test datum a Gaussian distribution, we use the KI. divergence between the Gaussians yielded by the approximate LLA method and vanilla LLA as proxy of the approximation error (averaged over set of test points). The results are reported in Appendix C.4. As shown, LLA-KFAC comes pretty close to LLA. Yet, LLA seems to underestimate in-between uncertainty in this setting, so ELLA seems to be a more reliable (instead more accurate) approximation than AFFAC. We also highlight the (see Figure 4(c)), which reflects that ELLA can strike
higher scalability of ELLA than good trade-off between efficacy and efficiency.
5.4 CIFAR-10 Classification
Then, we evaluate ELLA on CIFAR-10 benchmark using ResNet architectures [18]. We obtain pretrained MAP models from the open source community. Apart from MAP, LLA", LLA-Diag, LLA-KFAC, we further introduce last-layer LLA with KFAC approximation mean-field VI via Bayesian finetuning [10] (MFVI-BF) as baselines. LLA cannot directly applied huge. These methods all locate in the family of Gaussian approximate posteriors and are all post-hoc applied to the pretrained models, so the comparisons will be fair. Regarding the setups, use M 2000 and K 20 for ELLA:6 use 20 MC samples to estimate
-FFC) and
the GGN matrices
'Flipout [56] trick is deployed for variance-reduced gradient estimation.
*Storing 2000 vectors of size 1 costly, trade time for memory and compute them whenever needed.

---

0.2
TLLA
MAA
MEVLar
Mir
amnar
u4
BARRAC
Corfiisa Thacklle
(a)
(b)
(c)
Figure 4: (a) Error confidence plots for methods trained CIFAR-10 and tested CIFAR-I04SVHN,
(h) ELLLA varies w.r.t. CIFAR-10, (c) Comparison on wall-clock time used for predicting all CIFAR-10 test data (measured on an NVIDIA A40 GPU). The results are obtained with ResNet-20 architecture.
Table 2: Comparison test accuracy (%) f. NLL J, and ECB & on ImageNet. We report the average results over 3 random runs.
ResNet-18 NLL
ResNet-34 NLL
ResNet-50 NLL
Method
Acc.
ECE
Acc.
ECE
Acc.
ECE
ELLA
69.8
1.243
0.015
73.3
1.072
0.018
76.2
0.948
0.018
MAP
69.8
1.247
0.026
73.3
1.081
0.035
76.2
0.962
0.037
MFVI-BF
70.3
1.218
0.042
73.7
1.043
0.033
76.1
0.945
0.030
the posterior predictive of MFVI-BF (as incurs 20 NN forward passes), and use 512 ones for the other methods stated. We have enabled the tuning of the prior precision for all LLA baselines, but not for ELLA.
We present the comparison on test accuracy, NLL, and ECE in Table 1. As shown, ELLA exhibits superior NLL and ECE across settings. MFVI-BF also gives good NLL. LLA" and LLA* KFAC improve the uncertainty and calibration of MAP, yet underperforming ELLA. LLA-Do --iag and LLA-KFAC fail for unclear reasons (also reported by [9]), we thus exclude them from the following studies.
We then examine the models on the widely used benchmark CIFAR-10 corruptions [19] and report the results in Figure 3 and Appendix C.5. It is prominent that ELLA surpasses baselines in aspects of NLL and ECE at various levels skew, implying its ability to make conservative predictions for OOD inputs.
(OOD) generalization/rouutnees
We further evaluate the models on the combination of the CIFAR-10 test data and the OOD SVHN test data. The predictions on SVHN are all regarded as wrong due to label shift. We plot the average error rate of the samples with < 7 (7 [0, 1)) confidence in Figure (a). As shown, ELLA makes less mistakes than the basclines under various confidence thresholds. Figure 4 (b) displays how K impacts the test NLL. We see that K [20, 30) can already lead to good performance, reflecting the efficiency of ELLA. Another benefit of ELLA is that with it, we can actively control the performance vs. cost trade-off. Figure (c) shows the comparison on the time used for predicting all CIFAR-10 test data. We note that ELLA is slightly slower than MFVI and substantially faster than non-last-layer LLA methods.
5.5 ImageNet Classification
We apply ELLA to ImageNet classification [8] to demonstrate its scalability. The experiment settings identical with those for CIFAR-10. We observe that all L.AA methods implemented with Laplace would cause out-of-memory (OOM) errors or suffer from very long fitting procedures, thus take only MAP and MFVI-BF as baselines. Table 2 presents the comparison test accuracy, NLL,, and ECE with ResNet architectures. We see that ELLA maintains its superiority in ECE while MFVI-BF can induce higher accuracy and lower NLL. This may be attributed to that the pretrained MAP model lies at a sharp maxima of the true posterior so it is necessary to properly adjust the mean of the Gaussian approximate posterior.

---

MAP
Skew
Figure 5: NLL (Left) ECE (Right) on ImageNet corruptions for models trained with ViT-B architecture. Each box corresponds to summary of the results across types of skew.
We lastly apply ELLA to ViT-B [11]. We compare ELLA to MAP in Table 3 as all other baselines incur OOM errors or crushingly long running time. As shown, ELLA beats MAP in multiple as-
Table Comparison on test accu-
racy (%) f. NLL J. and ECE 4 on
ImageNet with ViT-B architecture.
pects. Figure shows the results of ELLA and MAP on ImageNet corruptions [19]. They are consistent with those for CIFAR-10 corruptions. We reveal by this experiment that ELLA can be more applicable and scalable method than most of existing BNNs.
Method
Acc.
NLL
ECE
ELLA
81.6
0.695
0.022
MAP
81.5
0.700
0.039
6 Related Work
LA [38, 50, 14, 30, 22, 6, 7] locally approximates the Bayesian posterior with a Gaussian distribution, analogous VI with Gaussian variationals [16, 1, 36, 52, 46) and SWAG [39]. LA can be applied to pretrained models effortlessly while the acquired posterions are potentially restrictive (see Section 5.5). VI enjoys higher flexibility yet relies costly training: Bayest [10] seems be remedy to the issue but its accessibility is still lower than LA. SWAG stores series of SGD iterates to heuristically construct an approximate posterior and is empirically weaker than LA/LLA [6].
Though more expressive approaches like deep ensembles [32] and MCMC [55, 4, 63] can explicitly explore diverse posterior modes, they face limitations in efficiency and scalability. What's 's more, it has been shown that LA/LLA can perform on par with or better than deep ensembles and cyclical MCMC on multiple benchmarks [6]. This may be attributed to the un ideeidetiiidd complicated relationships between the parameters space and the function space of DNNs [59]. Also of note that LA can embrace deep ensembles to capture multiple posterior modes [13].
[9] introduces general kernel approximation technique using neural networks. By contrast, we focus on leveraging kernel approximation to accelerate Laplace approximation. Thus, the focus of the two works is distinct. Indeed, this work shares a similar idea with the Sec 4.3 in [9] that LLA can be accelerated by kernel approximation. But, except for such idea, this work differentiates from [9] in aspects like motivations, techniques, implementations, theoretical backgrounds, and applications. Our implementation, theoretical analysis, and some empirical findings all novel.
Conclustion
This paper proposes ELLA, simple, effective, and reliable approach for Bayesian deep learning. ELLA addresses the unreliability issues of existing approximations to LLA and is implemented based on Nyström method. We offer theoretical guarantees for ELLA and perform extensive studies to testify its efficacy and scalability. ELLA currently accounts for only the predictive, and extending it estimate model evidence for model selection [38] deserves future investigation. Using model evidence to select the number of remaining eigenpairs K for ELLA is also viable.
Acknowledgmente and Disclosure of Funding
This work was supported by the National Key Research and Development Program of China (No. 2017YFA0700904) NSF of China Projects (Nos. 62061136001, 62076145, 62106121, U19B2034, U181 1461, U19A2081, 6197222); Beijing NSF Project (No. JQ19016); grant from Tsinghua Institute for Guo Qiang: and the High Performance Computing Center, Tsinghua University. J.Z. was also supported by the XPlorer Prize
10

---

References
[1] Charles Blundell, Julien Cornebise, Koray Kavukeuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International Conference Machine Learning, pages 1613-1622, 2015.
[2] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Millne and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
[3] David Burt, Carl Edward Rasmussen, and Mark Van Der Wilk. Rates of convergence for sparse variational gaussian process regression. In International Conference on Machine Learning, pages 862-871. PMLR, 2019.
[4] Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In International Conference on Machine Learning, pages 1683 1691. PMLR, 2014.
[5] Corinna Cortes, Mehryar Mohri, and Ameet Talwalkar, On the impact of kemel approximation on learning accuracy. In International Conference on Artificial Intelligence and Statistics, pages 113-120. JMLR Workshop and Conference Proceedings. 2010.
[6] Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace redux-effortless bayesian deep learning. Advances in Neural Information Processing Systems, 34, 2021.
[7] Erik Daxberger. Eric Nalisnick, James U Allingham, Javier Antorán, and José Miguel Hernández- Bayesian deep learning via subnetwork inference. In International Conference Machine Learning. pages 2510-2521, PMLR, 2021.
[8] J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li. and L. Fei-Fei. ImageNet: A Large-Scale Hierarchical Image Database. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2009.
[9] Zhijie Deng, Jiaxin Shi, and Jun Zhu. Neuralef: Deconstructing kernels by deep neural networks. arXiv preprint arXiv:2205.001
2022.
[10] Zhijie Deng, Hao Zhang, Xiao Yang. Yinpeng Dong. and Jun Zhu. Bayesadapter: Being bayesian. inexpensively and reliably, via bayesian tine-tuning. arXiv preprint arXiv:2010.01999, 2020.
[11] Alexey Dosovitskiy. Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn. Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold. Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representattonss 2020.
[12] Petros Drineas, Michael W Mahoney, and Nello Cristianini. On the nyström method for approximating a gram matrix for improved kernel-based learning. Journal of Machine Learning Research, 6(12), 2005.
[13] Runa Eschenhagen, Erik Daxberger. Philipp Hennig, and Agustinus Kristiadi. Mixtures of laplace approximations for improved post-hoc uncertainty in deep learning. arXiv preprint arXiv:2 2021.
[14] Andrew YK Foong, Yingzhen Li. José Miguel Hernández-Lobato and Richard E Turner 'in- between' uncertainty in bayesian neural networks. arXiv preprint arXiv: /906. /1537, 2019.
[15] Deena Francis and Kumudha Raimond. Major advancements kernel function approximation. Artificial Intelligence Review, 54(2):843-876, 2021.
[16] Alex Graves. Practical variational inference for neural networks. Advances in Neural Information Processing Systems, pages 2348-2356, 2011.
[17] Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. arXiv preprint arXiv: 045999 2017.
[18] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pages 770-778, 2016.
[19] Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. arXiv preprint arXiv: /903. /2261, 2019.
[20] José Miguel Hrrnánddez-Lobatt and Ryan Adams. Probabilistic backpropagation for scalable learning of Bayesian neural networks. In International Conference on Machine Learning, pages 1861-1869, 2015.
[21] Geoffrey Hinton and Drew Van Camp. Keeping neural networks simple by minimizing the description length of the weights. In ACM Conference on Computational Learning Theory, 1993.
[22] Alexander Immer, Maciej Korzepa, and Matthias Bauer. Improving predictions of bayesian neural nets via local linearization. In International Conference on Artificial Intelligence and Statistics, pages 703-711 PMLR, 2021.
11

---

in adam. In International Conference on [23] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning, pages 448-456. PMLR, 2015.
[24] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kemel:: Convergence and generalization in neural networks. arXiv preprint arXiv: 07572, 2018.
[25] Huisu Jang and Jaewook Lee. An empirical study on modeling and prediction of bitcoin prices with bayesian neural networks based on blockchain information. leee Access, 6:5427-5437, 2017.
[26] Rong Jin, Tianbao Yang, Mehrdad Mahdavi, Yu-Feng Li, and Zhi-Hua Zhou. Improved bounds for the nyström method with application to kernel classification. IEEE Trunsactions on Information Theory. 59(10))6939-69999 2013.
[27] Alex Kendall and Yarin Gal. What uncertainties do we need in Bayesian deep learning for computer vision? In Advances in Neural Information Processing Systems, pages 5574-5584, 2017.
[28] Mohammad Emtiyaz. Khan, Alexander Immer, Ehsan Abedi, and Maciej Korzepa. Approximate inference turns deep networks into gaussian processes. arXiv preprint arXiv: 10:::9906000 00660/930, 2019.
[29] Mohammad Emtiyaz Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivastava. Fast and scalable Bayesian deep learning by Machine Learning, pages 2616-2625, 2018.
[30] Agustinus Kristiadi, Matthias Hcin, and Philipp Hennig. Being bayesian, even just a bit, fixes overconfi- dence in relu networks. arXiv preprint arXiv:2002. 10/18, 2020.
[31] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[32] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Advances in Neural Information Processing Systems, pages 6402-64006 6413, 2017,2017.
[33] Neil David Lawrence. Variational inference in probabilistic models. PhD thesis, Citeseer, 2001.
[34] Christian Leibig, Vaneeda Allken, Murat Seçkin Ayhan, Philipp Berens, and Siegfried Wahl. Leveraging uncertainty information from deep neural networks for disease detection. Scientific Reports, 7(1):1114, 2017.
[35] Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose Bayesian inference algorithm. In Advances in Neural Information Processing Systems, pages 2378- 22766, 2016.
[36] Christos Louizos and Max Welling. Structured and efficient variational deep learning with matrix gaussian posteriors, In International Conference on Machine Learning, pages 1708-1716, 2016.
[37] David JC MacKay. A practical Bayesian framework for hackpropagation networks. Newral Computation, 4(3):448-72, 1992.
[38] David John Cameron Mackay. Bayesian methods for adaptive models. PhD thesis, California Institute of Technology, 1992.
[39] Wesley Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrow, and Andrew Gordon Wilson. A simple baseline for bayesian uncertainty in deep learning. Advances in Neural Information Processing Systems, pages 13153-13164, 2019.
[40] James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In International Conference on Machine Learning, pages -408-241.. PMLR. 2015.
[41] Andres Masegosa. Learning under model misspecification: Applications variational and ensemble methods. Advances Neural Information Processing Systems, 33:5479-5491, 2020.
[42] Mercer. Functions ofpositive and negativetypeand theircommection with the thcory ofintegral equations. Philosophicol Trinsdictions of the Rogyal Society, pages 4-415, 1909.
[43] Marina Munkhoeva, Yermek Kapushev, Evgeny Burnaev, and Ivan Oseledets. Quadrature-based features for kemel approximation. Advances Neural Information Processing Systems, 31, 2018.
[44] Radford M Neal. Bayesian Learning for Neural Networks. PhD thesis, University of Toronto, 1995.
[45] Evert Nyström. Über die praktische auflösung von integralgleichungen mit anwendungen auf randwer- taufgaben. Acta Mathematica, 54:185-204, 1930.
[46] Kazuki Osawa, Siddharth Swaroop. Anirudh Jain, Runa Eschenhagen, Richard E Turner, Rio Yokota, and Mohammad Emtiyaz Khan. Practical deep learning with Bayesian principles. arXiv preprint arXiv: (90.. 0200600206. 2019.
[47] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, perf-reeformance deep learning library. In Advances in Neural Information Processing Systems, pages 8026 8037, 2019.
12

---

[48] Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. Advances in Neural Information Processing Systems, 20, 2007.
[49] Ali Rahimi and Benjamin Recht. Weighted sums of random kitchen sinks: Replacing minimization with randomization in learning, Advances in Neural Information Processing Systems, 21, 2008.
[50] Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable laplace approximation for neural networks. In 6th International Conference on Learning Representations, volume 6. International Conference on Representation Learning, 2018.
[51] Matthias Seeger. Gaussian processes for machine learning. International journal of neural systems, 14(02))69-106, 2004.
[52] Shengyang Sun, Changyou Chen, and Lawrence Carin. Learning structured weight uncertainty in Bayesian neural networks. In International Conference Artificial Intelligence and Statistics, pages 1283-1292, 2017.
[53] Michalis Titsias. Variational learning of inducing variables in sparse gaussian processes. In International Conference on Artificial Intelligence and Statistics, pages 567-574, PMLR, 2009,
[54] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, pages 5998-600S 2017.
[55] Max Welling and Yee W Tch. Bayesian learning via stochastic gradient langevin dynamics. International Conference on Machine Learning, pages 681 681-688, 2011.
[56] Yeming Wen, Paul Vicol, Jimmy Ba, Dustin Tran, and Roger Grosse. Flipout: Efficient pseudo-independent weight perturbations on mini-batches. arXiv preprint 1803,v:800004430 2018.
[57] Veit Wild, Motonobu Kanagawa, and Dino Sejdinovic. Connections and equivalences between the nystr'." om method and sparse variational gaussian processes. arXiv preprint arXiv:2106.01122 2021.
[58] Christopher Williams and Matthias Seeger. Using the nyström method to speed up kernel machines.
Advances in Neural Information Processing Systems, 13, 2000.
[59] Andrew Gordon Wilson and Pavel Izmailov. Bayesian deep learning and a probabilistic perspective of generalization. arXiv preprint arXiv:2002.0879. 2020.
[60] Max A Woodbury. Inverting modified matrices. Statistical Research Group, 1950.
[61]
Felix Xinnan X Yu, Ananda Theertha Suresh, Krzysztof M Choromanski, Daniel N Holtmann-Riee -iccn, and Sanjiv Kumar. Orthogonal random features. Advances in Neural Information Processing Systems, 29, 2016.
[62] Guodong Zhang, Shengyang Sun, David Duvenaud, and Roger Grosse. Noisy natural gradient variational inference. In International Conference on Machine Learning, pages 5847-5856, 2018.
[63] Ruqi Zhang, Chunyuan Li, Jianyi Zhang, Changyou Chen, and Andrew Gordon Wilson. Cyclical stochastic gradient meme for bayesian deep learning. arXiv preprint arXiv:1902.03933 2019.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Section Related Work.
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [Yes] See Appendix.
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)? [Yes] See supplemen- tal material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Section Experiments.
13

---

(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [No]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, cloud provider)? [No]
4. If you are using existing assets (e.g., code, data, models) or curating/releasig new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material URL? [No]
(d) Did you discuss whether and how consent obtained from people whose data you' 're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "g3lUD2uUTKT", "reviewer": "Reviewer_r6Sm", "summary": "This work proposes a new, scalable variant of the (linearized) Laplace approximation (LLA) for inferring posterior distributions over the weights of deep learning models. Leveraging connections between LLA and neural tangent kernels (NTKs) -- namely that the LLA can essentially be viewed in function space as a Gaussian process with the NTK -- the paper describes how to speed up the LLA via a Nyström approximation to the NTK, effectively yielding an implicit sparse approximation to the LLA-GP predictive distribution. The work then elaborates on how to efficiently implement the proposed approach using automatic differentiation frameworks, and provides a theoretical analysis of the induced approximation error for the kernel. The resulting ELLA method is then empirically shown to outperform several baselines (both other LLA variants and a variational method) across a range of benchmarks -- even large-scale results with a Vision Transformer architecture on ImageNet are reported. It is also empirically found that ELLA (and LLA methods more generally) tends to \"overfit\" in some settings, which is proposed to be remedied by only using a subset of the data for estimating the covariance matrix, similar in spirit to early stopping methods often used in standard neural network training.", "questions": "- l. 234-235: in the footnote, you mention that \"storing 2000 vectors of size $P$ is costly\", but do you not just need to store $K=20$ instead of $M=2000$ vectors of size $P$ here (cf. l. 163-164)?\n- l. 239-240: it would be good to understand why LLA-Diag and LLA-KFAC fail in this setting, given that previous works have been able to make them work fairly well, e.g. Daxberger et al. (2021); what have you tried to make these work and what is your intuition on why they fail?\n", "limitations": "Yes, the authors have adequately addressed this.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Summary of Review**\n\nThis work proposes a novel, theoretically-justified and clever variant of the Laplace approximation that is empirically shown to be performant and scalable to fairly modern deep learning settings, and that could therefore have significant impact by enabling practical uncertainty quantification for deep neural networks. Overall, there are many things that I liked about this paper; however, I also see some significant issues with the empirical evaluation (see details below). All in all, I am inclined to recommend acceptance of the manuscript, but would be more enthusiastic in my judgement if the authors can convincingly address the concerns raised.\n\n**Strengths**\n- The paper is overall well-written, easy-to-follow, and clearly-structured.\n- The proposed ELLA method for accelerating LLA is novel, theoretically sound, and intuitively sensible; it is therefore a timely and welcome addition to the fast-growing Laplace toolbox for Bayesian deep learning.\n- I appreciate the implementation-specific details in Section 3.4 including the code snippets in Algorithm 1-2, which make it fairly clear how one would practically implement ELLA (without needing to look at the actual Python code).\n- The theoretical results on the approximation error induced by ELLA are nice-to-have (even though they appear to be somewhat straightforwardly inherited from previous analyses of Nyström-like kernel approximations).\n- A main advantage of the proposed ELLA method is that a practitioner can actively control the performance vs. cost trade-off by choosing the hyperparameters $M$ and $K$, which is in contrast to most other methods (e.g. LLA-Diag and LLA-KFAC) that have a fixed cost and performance (however, as mentioned in the weaknesses below, this trade-off should be assessed more thoroughly in the empirical evaluation).\n- The empirical evaluation demonstrates that ELLA can outperform relevant baselines on image-classification-based uncertainty calibration and out-of-distribution detection tasks; I particularly appreciate the inclusion of experiments on ImageNet, which is (unfortunately) still not so common in the Bayesian deep learning literature; I was especially impressed by the results on the recent Vision Transformer architecture, which promisingly demonstrates that the method can be applied in fairly modern deep learning settings (which, again, is not typically the case -- Bayesian methods normally lack behind significantly when it comes to adoption of advances in deep learning).\n- I really liked the experiment showing that ELLA (and LLA methods more generally) have a tendency to \"overfit\" to larger datasets (in the sense that the test loss starts increasing again past a certain number of data points used for fitting the covariance matrix), demonstrating that LLA methods can be significantly sped-up by subsampling the data in some settings.\n\n**Weaknesses**\n- As all experiments seem to have been run with just a single random seed, the results do not come with any error bars, which makes it difficult to reason about the statistical significance of the reported conclusions; while I understand that repeating experiments for multiple random seeds linearly increases the required computational effort, reporting some sort of error statistic would significantly aid the credibility of the drawn conclusions; for the particularly expensive experiments (i.e. on ImageNet) I could perhaps accept lack of compute as an excuse (although even there, a minimum number of 3 seeds would be desirable), but on the cheaper experiments (i.e. on CIFAR-10) I would really expect some repetition of the experiments.\n- The empirical comparison does not seem entirely fair, as different methods have different memory and compute requirements; in particular, for the hyperparameters chosen, ELLA seems to be more expensive than some of the other baselines considered (in terms of compute and/or memory cost, as e.g. also shown in Fig. 4c); therefore, it might be more appropriate and insightful to plot performances in 2D with memory and/or compute effort on the x-axis and performance on the y-axis, so that it becomes clear how the different methods can trade-off cost vs. performance; for ELLA, it would be great to then report performance for different values of $M$ and/or $K$ (i.e. different costs), yielding a Pareto curve; a practitioner can then choose the method/hyperparameters on the Pareto front that best fulfills their needs, i.e. which is either 1) as cheap as possible for a given desired performance X, or 2) as performant as possible for a given desired cost Y; without such a plot it is difficult to draw conclusions on how superior ELLA really is -- although ELLA's ability to cater to a wide range of performance/cost trade-offs by sweeping over $M$ and $K$ should already make it look better than most of the baselines; you already show several plots that are related in spirit (Fig 2. a-b, Fig. 4 b-c), but I believe the specific kind of plot I described would be much more insightful.\n- For the ImageNet experiments, I would expect that at least the last-layer KFAC variant should be feasible to run and not yield an OOM error (with either a ResNet or ViT) -- e.g. the Eschenhagen et al. (2021) paper you cite seems to have managed to run this Laplace variant on ImageNet with a ResNet; I think it would significantly strengthen the paper to at least have this one Laplace baseline to compare with; if this is really an issue with the Laplace library, I would encourage you to reach out to the authors / raise an issue on the GitHub repo -- I also had issues with this library in the past, and in my experience, the authors are typically happy to help and fairly quick to respond.\n- The proposed approach can be viewed as limited in novelty in the sense that it effectively combines the LLA with a Nyström approximation to the NTK, which both have been studied (fairly extensively) before. That being said, I think the idea is pretty neat (even though somewhat obvious in hindsight) and appreciate that the authors made the effort of actually getting this to work, so I don't view this to be a major issue.\n\n**Minor issues**\n- l. 226-227: to back up the claim that \"ELLA delivers a closer approximation to LLA\" it might be worth explicitly quantifying the discrepancies between the predictive distributions using some metric; qualitatively (i.e. by just looking at the plots), it seems like LLA-KFAC also comes pretty close (but I agree that even LLA might underestimate in-between uncertainty in this setting, so the ELLA fit could be more desirable).\n- Generally, it would be useful to have short take-away messages in the captions of each table/figure for clarity and convenience.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "X8pYDnt3FMa", "reviewer": "Reviewer_2eHY", "summary": "The paper proposes ELLA, a Nyström approximation, as an alternative posterior (predictive) approximation for linearized Laplace. Their method uses the NTK-based formulation of linearized Laplace and improves over other common Laplace approximations in terms of quality of the predictive while maintaining a relatively low computational complexity.", "questions": "- ll 31 says that LLA has a more benign covariance. Can you elaborate on what that means or in what way it is benign?\n- Figure 1: how comes that ELLA gives better \"in-between\" uncertainty? Are the hyperparameters equivalent to the LLA method? Generally, ELLA approximates the LLA variant so it is unclear why/how it is better. This would be worth investigating more. \n- Could you elaborate on the simplification in Equation (6), for example, by putting a detailed derivation in the appendix? It is not clear to me where the equivalence comes from.\n- Regarding computational overhead (ll. 162 ff): Can you relate this to the prevelant alternative of KFAC-LLA? It would be interesting to know the complexity of computing the posterior predictive for an unseen data point after inference.\n- The paper at hand only discusses the posterior predictive. Have the authors considered using the Nyström approximation for marginal likelihoods as well. In [*scalable marginal likelihood estimation for model selection in deep learning*](https://arxiv.org/abs/2104.04975), the NTK variant for linearized Laplace was used but without further approximations. It would be interesting to experiment if Nyström can provide an alternative scalable method. Potentially, this could alleviate overfitting in Figure 2c) by selecting the number of bases $K$?\n- Figure 2 is hard to parse since the colorbars are unclear. What is the differene between Nyström and ELLA there in particular?\n- Table 2 has no error bars but bold figures. Which policy is used to make numbers bold? Did the authors run for several seeds? Same Q for Table 3.\n- Regarding the overfitting issue of ELLA: to what extent can changing the prior fix such overfitting issue? For example, the prior precision is often changed after training to improve the performance of LA [6]. Is this an option here, too?", "limitations": "There are no immediate societal impacts of the present work. Computational limitations are discussed in the paper, for example, the authors discuss the issue of overfitting present in their method.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n- Can reduce the cubic complexity in data points $N$ or parameters $P$ of linearized Laplace in neural networks down to much lower complexity that can be controlled by the number of samples used for Nyström. This way, a performance-accuracy trade-off can be made.\n- Performance of the proposed posterior predictive improves over cheap last-layer Laplace and is better than the full-network variants.\n- Interesting and open discussion of the overfitting issue of the proposed method ELLA. \n- Exhaustive experiments with ablations on the number of factors $K$ and several metrics of interest. Overall convinces that the proposed method is a valuable addition to the family of Laplace approximations.\n\n## Weaknesses\n- Only posterior predictive is discussed but not marginal likelihood estimation. Since the proposed method is rather simple (\"take linearized Laplace in its kernel variant and apply Nyström, which is well-known from GP literature\"), it would greatly strengthen the paper to at least show not only predictives but also the marginal likelihood, which should in principle be easy to add.\n- Misunderstanding of relationship between linearized Laplace and GGN. ll 3-6 and 30-32: linearized Laplace and Laplace with GGN are equivalent as shown in reference [22]. The equivalent variant with (neural tangent) kernels has been used before as well in [22,28] but without Nyström approximations. ll 92-99 discusses the scalability issues only of the parametric LLA which requires approximations to the $P \\times P$ GGN. The problem that is tackled in the present paper, however, is the approximation of the equivalent $N \\times N$ kernel, which is for example defined. Clearly, $N \\times N$ is still intractable so the proposed Nyström approximation is well justified.\n- ELLA does not seem to approximate full LA, which it should in theory. Figure 1: discrepancy of ELLA and LLA, it seems like ELLA does not get very close to LLA in Figure 1 and there is no explanation as to why that would be. In line 50, it is said though that the approximation becomes accurate.\n- Relation to prior work could be pointed out more clearly. For example, it is unclear to what extent the Theorems are novel/original or if they follow in a straightforward way from previous theoretical contributions. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "C0MAzEtEv0L", "reviewer": "Reviewer_AKXi", "summary": "The paper targets the improvement of probabilistic predictions in Bayesian deep learning with the (linearized) Laplace approximation. In particular, it proposes to apply the Nystrom kernel approximation to the neural-tangent-like covariance matrix of the corresponding Gaussian process that is the function-space dual of the linearized Laplace in weight-space. Using the Nystrom approximation allows the authors to scale the linearized Laplace to ImageNet and ViT scale. They compare their method to other Laplace approximations in weight space and find that their method performs better, though in some cases only marginally, at a similar or lower computational cost.\n", "questions": "* I found it hard to understand/follow the discussion of the overfitting issue/behaviour of ELLA in Sec 5.2. As far as I understood, ELLA uses pre-trained models and only deals with the GP predictive distribution. As such, \"early stopping\" does not really make sense to me as it typically refers to early stopping of training. Could you please elaborate?", "limitations": "* see Weaknesses above, in particular, the authors only investigate ELLA for predictions.\n* several works (e.g. ref [22], [6]) point out that the choice of prior precision is important to achieve good predictive performance (in particular for uncertainties/NLL/ECE); however, it looks like the values here have simply been taken from the pre-trained model regularisation, which may explain the poor performance of LLA and LLA+KFAC. For this reason, [6] (as far as I understand) proposes to use last-layer Laplace by default when these hyperparameters are not tuned/the MAP model itself does not overfit.\n* the illustrative regression example (sec 5.3) is not discussed in enough detail. In particular, it should be investigated and further discussed why ELLA can be better than the full GGN model (LLA), to which it should be equivalent (except for the Nystrom approximation, of course). The authors are \"surprised\" (see caption to Fig 1) by ELLA performing better than the full model (LLA) as well, yet they do not investigate or discuss this. For example, the full GP predictive (without Nystrom approximation) as well as different/larger values for M and K should be considered in this case to build further intuition for the model and to investigate this behaviour. ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* The Laplace approximation for larger scale machine learning has recently seen increasing attention and this paper further improves scalability by introducing the Nystrom approximation in its dual (function space) formulation.\n* The paper seems to be technically correct, though I haven't checked all the details.\n\nWeaknesses:\n* Generality + Novelty: The novelty and originality presented in this paper is limited, especially when compared to other papers at ICML or NeurIPS in this space (e.g. ref [6] provides a library for Laplace approximation and a thorough review/benchmark, or ref [9] introduces kernel approximations more generally and not only for large-scale Laplace predictions). While this paper extends the Laplace approximation to ImageNet and ViTs, the paper ultimately discusses one particular approximation (Nystrom) to one particular covariance function (NTK) for one particular task (probabilistic predictions on pre-trained models); it largely builds on existing libraries and improvements over MAP predictions are often marginal (especially on some of the larger-scale problems) and provided without errorbars. The authors also claim that no other methods can be applied to ViT or ImageNet, yet last-layer Laplace (LLA*) should be applicable in this case as well (at least when subsampling the dataset; e.g. [6] apply last layer Laplace to transformer models).\n* Contributions + relation to prior work: \n  - Several results in Sec 3.2 and 3.3 are, in my opinion, not novel and have been similarly derived in the kernel literature (where the Nystrom approximation is commonplace) as well as in Appendix A.2 and A.3 of NeuralEF (ICML2022, reference [9] in the paper under review) for the NTK kernel. While \"NeuralEF\" is referenced (as ref [9]) in the paper, it is not mentioned in Sec 3.2 and 3.3 of this paper and it is not clarified, which results in Sec 3.2 and 3.3 are novel and which are already known.\n  - The Nystrom approximation is a well known and well studied approximation to kernel functions, yet a discussion/presentation of this is missing from this paper.\n  - NeuralEF [9] highlights several limitations of the Nystrom approximation when applied to NTK mostly to do with scalability (see Sec 2.2 in [9] in particular). They highlight the cost of the eigendecomposition as well as prediction costs at test time. I agree with [9] that this could be problems and scalability of the Nystrom should be discussed more prominently.\n* also see \"Limitations\" below\n* I found the presentation/clarity to be not clear:\n\t- The setting (probabilistic prediction with Bayesian neural networks), ultimate goal/objective (improve linearized(?) posterior predictive), as well as the baseline methods (discussion + equations for e.g. Laplace with KFAC, ...) should be introduced more clearly. E.g. the posterior predictive distributions  evaluated are only mentioned in passing/do not have clear equations/equation numbers\n\t- The paper ultimately presents a technical contribution (Nystrom for this particular kernel) in an arguably very technical way. In my opinion, more emphasis should be put on explanation and some of the technical details (e.g. indexing with $i$ in Sec 3.3) should be left to the appendix.\n\t- Alg 1 and 2 should be explained in much more detail; comments + references to the text as well as from the text to the algorithms are missing\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept